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/29] [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/29] [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/29] [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/29] [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/29] [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/29] [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/29] [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/29] [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/29] [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/29] [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 7333ac8e86a3d4b3fab8be3f335351c99c8768a0 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:32:22 -0700 Subject: [PATCH 11/29] [None][feat] MiniMax-M3 NVFP4 low latency decode megakernel for Rubin (sm_107) TP2 / TP4 / TP8 in modeling_v2 Three modeling_v2 targets for the MiniMax-M3 NVFP4 checkpoint on Rubin (sm_107), minimax_m3_nvfp4__sm_107__tp2, __tp4 and __tp8, in the flat layout of the sm_103 targets: one modeling.py per target, a weights.py weight table and an ops/ catalog entry with certified cells and a test for every kernel the model calls. Routing picks a target only for the matching checkpoint, SM and width. sm_107 enablement: Arch.SM_107 in the catalog and the (10, 7) routing entries; the compiled trtllm::minimax_m3_sm107* kernels (built for sm_107 only; torch ops in thop/minimaxM3/, fake impls in the M3 fake_ops module) and the shared TF32 router GEMV and spec-tail draft-input kernels; the CuTe DSL decode MoE engines, MX GEMV and fused embedding + first RMSNorm; the decode runtime pieces (step-input pack, graph runner hooks, step fold, MSA decode metadata reuse, the SM107 NVFP4 -> FP8 KV fallback skipping MiniMax-M3 MSA, late-compile detection), each inert until a target's runtime hooks turn it on; the FMHA fallback serving SM107 paged prefill at the TRT-LLM Gen extra page sizes (the 128-token KV pages its MiniMax-M3 caches opt into); and the scoped PrimTS work-tile rebind in the vendored FlashInfer patch. The TP2 target serves decode steps of 1..192 rows: context steps of 9..48 rows route on the tensor-core router, steps with context rows run under breakable prefill CUDA graphs (the drafter's attention in an eager break) and take the MC finalize boundary at 9..48 rows, and decode steps of >= 33 requests take the stock MSA merge path (MergeMxfp8.MAX_ROWS 128; no canon profile exceeds 96 rows). Runtime items: pinned FlashInfer CuTe DSL MXFP8 tactics above 32 rows on SM107, for models whose runtime hooks opt in (the MiniMax-M3 sm_107 targets do); the AllReduce tuner leaves NCCL_SYMMETRIC out of its candidates only when TRTLLM_AR_TUNE_NCCL_SYMMETRIC=0 (opt-in; the default 1 keeps the old candidates); PyExecutor emits a request's first content after two iterations when the model's runtime hooks ask for it (the MiniMax-M3 sm_107 targets do; other models keep the stock schedule). No lossy steps: every kernel keeps the stock op's precision. The TP4 and TP8 targets produce bitwise the same tokens and per-step outputs as the pre-restructure sm_107 targets, with and without Eagle3. Two known deltas against the previous wide TP2 target (the 9..128-row cells this commit folds in), both measured with greedy sampling: under Eagle3, at 12 or more concurrent requests the drafter's BF16 linears past its low-M tactic table can resolve a near-tie differently from the in-tree drafter, so some draft tokens differ, while the target's tokens are identical and its logits bitwise; without speculation (not a production configuration), decode steps of 9..48 rows are bitwise on the first such step and differ after it (cause not isolated), with identical output tokens in every comparison. Steps of 1..8 rows are bitwise to the previous TP2 target. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 235 +- 3rdparty/vendor_sources.lock.yaml | 4 +- cpp/tensorrt_llm/kernels/CMakeLists.txt | 32 +- .../kernels/minimaxM3/routerGemv.cu | 67 +- .../kernels/minimaxM3/routerGemv.h | 4 + .../kernels/minimaxM3/specTail.cu | 25 +- .../minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu | 267 + .../minimaxM3Sm107Tp2/denseQkNormRopeFp8.h | 64 + .../kernels/minimaxM3Sm107Tp2/draftStep.cu | 620 ++ .../kernels/minimaxM3Sm107Tp2/draftStep.h | 111 + .../minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu | 189 + .../minimaxM3Sm107Tp2/drafterRopeFp8Kv.h | 79 + .../kernels/minimaxM3Sm107Tp2/foldTail.cu | 86 + .../kernels/minimaxM3Sm107Tp2/foldTail.h | 72 + .../minimaxM3Sm107Tp2/fullCtxReduce.cu | 223 + .../kernels/minimaxM3Sm107Tp2/fullCtxReduce.h | 71 + .../minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu | 1092 +++ .../minimaxM3Sm107Tp2/mcAllreduceNormTp2.h | 83 + .../kernels/minimaxM3Sm107Tp2/producerPw2.cu | 339 + .../kernels/minimaxM3Sm107Tp2/producerPw2.h | 76 + .../minimaxM3Sm107Tp2/sampledAccept.cu | 596 ++ .../kernels/minimaxM3Sm107Tp2/sampledAccept.h | 82 + .../kernels/minimaxM3Sm107Tp2/selectLong2.cu | 246 + .../kernels/minimaxM3Sm107Tp2/selectLong2.h | 50 + .../kernels/minimaxM3Sm107Tp2/siluAndMul.cu | 89 + .../kernels/minimaxM3Sm107Tp2/siluAndMul.h | 40 + .../minimaxM3Sm107Tp2/sparseDecodeLc2.cu | 640 ++ .../minimaxM3Sm107Tp2/sparseDecodeLc2.h | 77 + .../kernels/minimaxM3Sm107Tp2/specOffsets.cu | 58 + .../kernels/minimaxM3Sm107Tp2/specOffsets.h | 41 + .../kernels/minimaxM3Sm107Tp2/stepPrologue.cu | 111 + .../kernels/minimaxM3Sm107Tp2/stepPrologue.h | 47 + .../kernels/minimaxM3Sm107Tp2/swigluOai.cu | 100 + .../kernels/minimaxM3Sm107Tp2/swigluOai.h | 43 + .../minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu | 296 + .../minimaxM3Sm107Tp4/denseQkNormRopeFp8.h | 68 + .../kernels/minimaxM3Sm107Tp4/draftStep.cu | 667 ++ .../kernels/minimaxM3Sm107Tp4/draftStep.h | 107 + .../minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu | 189 + .../minimaxM3Sm107Tp4/drafterRopeFp8Kv.h | 79 + .../kernels/minimaxM3Sm107Tp4/foldTail.cu | 64 + .../kernels/minimaxM3Sm107Tp4/foldTail.h | 129 + .../minimaxM3Sm107Tp4/fullCtxReduce.cu | 268 + .../kernels/minimaxM3Sm107Tp4/fullCtxReduce.h | 77 + .../minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu | 1264 ++++ .../minimaxM3Sm107Tp4/mcAllreduceNormTp4.h | 83 + .../kernels/minimaxM3Sm107Tp4/producerPw4.cu | 340 + .../kernels/minimaxM3Sm107Tp4/producerPw4.h | 76 + .../minimaxM3Sm107Tp4/sampledAccept.cu | 944 +++ .../kernels/minimaxM3Sm107Tp4/sampledAccept.h | 87 + .../kernels/minimaxM3Sm107Tp4/selectLong2.cu | 312 + .../kernels/minimaxM3Sm107Tp4/selectLong2.h | 52 + .../kernels/minimaxM3Sm107Tp4/siluAndMul.cu | 89 + .../kernels/minimaxM3Sm107Tp4/siluAndMul.h | 40 + .../minimaxM3Sm107Tp4/sparseDecodeLc4.cu | 743 ++ .../minimaxM3Sm107Tp4/sparseDecodeLc4.h | 104 + .../kernels/minimaxM3Sm107Tp4/specOffsets.cu | 58 + .../kernels/minimaxM3Sm107Tp4/specOffsets.h | 41 + .../kernels/minimaxM3Sm107Tp4/stepPrologue.cu | 111 + .../kernels/minimaxM3Sm107Tp4/stepPrologue.h | 47 + .../kernels/minimaxM3Sm107Tp4/swigluOai.cu | 100 + .../kernels/minimaxM3Sm107Tp4/swigluOai.h | 43 + .../kernels/minimaxM3Sm107Tp8/draftStep.cu | 711 ++ .../kernels/minimaxM3Sm107Tp8/draftStep.h | 119 + .../minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu | 189 + .../minimaxM3Sm107Tp8/drafterRopeFp8Kv.h | 79 + .../kernels/minimaxM3Sm107Tp8/foldTail.cu | 64 + .../kernels/minimaxM3Sm107Tp8/foldTail.h | 129 + .../minimaxM3Sm107Tp8/fullCtxReduce.cu | 268 + .../kernels/minimaxM3Sm107Tp8/fullCtxReduce.h | 77 + .../minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu | 1096 +++ .../minimaxM3Sm107Tp8/mcAllreduceNormTp8.h | 83 + .../kernels/minimaxM3Sm107Tp8/producerPw8.cu | 340 + .../kernels/minimaxM3Sm107Tp8/producerPw8.h | 76 + .../minimaxM3Sm107Tp8/sampledAccept.cu | 946 +++ .../kernels/minimaxM3Sm107Tp8/sampledAccept.h | 87 + .../kernels/minimaxM3Sm107Tp8/siluAndMul.cu | 89 + .../kernels/minimaxM3Sm107Tp8/siluAndMul.h | 40 + .../minimaxM3Sm107Tp8/sparseDecodeLc8m.cu | 746 ++ .../minimaxM3Sm107Tp8/sparseDecodeLc8m.h | 105 + .../kernels/minimaxM3Sm107Tp8/specOffsets.cu | 58 + .../kernels/minimaxM3Sm107Tp8/specOffsets.h | 41 + .../kernels/minimaxM3Sm107Tp8/stepPrologue.cu | 111 + .../kernels/minimaxM3Sm107Tp8/stepPrologue.h | 47 + .../kernels/minimaxM3Sm107Tp8/swigluOai.cu | 100 + .../kernels/minimaxM3Sm107Tp8/swigluOai.h | 43 + cpp/tensorrt_llm/kernels/quantization.cu | 30 + cpp/tensorrt_llm/kernels/quantization.cuh | 58 + .../trtllmGenKernels/fmha/fmhaKernels.h | 129 +- .../thop/minimaxM3/minimaxM3RouterGemvOp.cpp | 23 +- .../minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp | 125 + .../minimaxM3Sm107Tp2DraftStepOp.cpp | 220 + .../minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp | 103 + .../minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp | 130 + .../minimaxM3Sm107Tp2FullCtxReduceOp.cpp | 74 + .../minimaxM3Sm107Tp2McAllreduceNormOp.cpp | 181 + .../minimaxM3Sm107Tp2ProducerPwOp.cpp | 185 + .../minimaxM3Sm107Tp2SampledAcceptOp.cpp | 110 + .../minimaxM3Sm107Tp2SelectLongOp.cpp | 85 + .../minimaxM3Sm107Tp2SiluAndMulOp.cpp | 66 + .../minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp | 169 + .../minimaxM3Sm107Tp2SpecOffsetsOp.cpp | 74 + .../minimaxM3Sm107Tp2StepPrologueOp.cpp | 68 + .../minimaxM3Sm107Tp2SwigluOaiOp.cpp | 65 + .../minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp | 128 + .../minimaxM3Sm107Tp4DraftStepOp.cpp | 215 + .../minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp | 103 + .../minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp | 89 + .../minimaxM3Sm107Tp4FoldTailParams.h | 111 + .../minimaxM3Sm107Tp4FullCtxReduceOp.cpp | 92 + .../minimaxM3Sm107Tp4McAllreduceNormOp.cpp | 181 + .../minimaxM3Sm107Tp4ProducerPwOp.cpp | 185 + .../minimaxM3Sm107Tp4SampledAcceptOp.cpp | 113 + .../minimaxM3Sm107Tp4SelectLongOp.cpp | 87 + .../minimaxM3Sm107Tp4SiluAndMulOp.cpp | 66 + .../minimaxM3Sm107Tp4SparseDecodeLcOp.cpp | 191 + .../minimaxM3Sm107Tp4SpecOffsetsOp.cpp | 74 + .../minimaxM3Sm107Tp4StepPrologueOp.cpp | 68 + .../minimaxM3Sm107Tp4SwigluOaiOp.cpp | 65 + .../minimaxM3Sm107Tp8DraftStepOp.cpp | 236 + .../minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp | 103 + .../minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp | 89 + .../minimaxM3Sm107Tp8FoldTailParams.h | 111 + .../minimaxM3Sm107Tp8FullCtxReduceOp.cpp | 92 + .../minimaxM3Sm107Tp8McAllreduceNormOp.cpp | 181 + .../minimaxM3Sm107Tp8ProducerPwOp.cpp | 185 + .../minimaxM3Sm107Tp8SampledAcceptOp.cpp | 113 + .../minimaxM3Sm107Tp8SiluAndMulOp.cpp | 66 + .../minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp | 189 + .../minimaxM3Sm107Tp8SpecOffsetsOp.cpp | 74 + .../minimaxM3Sm107Tp8StepPrologueOp.cpp | 68 + .../minimaxM3Sm107Tp8SwigluOaiOp.cpp | 65 + .../_experimental/modeling_v2/catalog/_op.py | 1 + .../modeling_v2/models/minimax_m3/fake_ops.py | 82 + .../minimax_m3_nvfp4__sm_107__tp2/__init__.py | 3 + .../kernels/__init__.py | 3 + .../kernels/embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../kernels/embedding_firstnorm/op.py | 143 + .../kernels/long_select.py | 70 + .../kernels/m3_attn_lc.py | 164 + .../kernels/m3_draft_step.py | 341 + .../kernels/m3_drafter_attn.py | 413 ++ .../kernels/m3_fold_tail.py | 172 + .../kernels/m3_force_vec.py | 63 + .../kernels/m3_fullctx_decode.py | 296 + .../kernels/m3_index_score.py | 670 ++ .../kernels/m3_locality.py | 448 ++ .../kernels/m3_lowm_direct.py | 497 ++ .../kernels/m3_lowm_splitk.py | 1249 ++++ .../kernels/m3_mc_allreduce.py | 392 + .../kernels/m3_moe/__init__.py | 16 + .../kernels/m3_moe/m3_moe_kernel.py | 3936 ++++++++++ .../kernels/m3_moe/op.py | 808 +++ .../kernels/m3_moe_finalize.py | 95 + .../kernels/m3_prequant.py | 111 + .../kernels/m3_report.py | 34 + .../kernels/m3_router_gemv.py | 107 + .../kernels/m3_sampled_accept.py | 203 + .../kernels/m3_select_long_dsl.py | 366 + .../kernels/m3_shard_accept.py | 223 + .../kernels/m3_spec_tail.py | 113 + .../kernels/m3_target_head.py | 113 + .../kernels/m3_tokmap_guard.py | 185 + .../kernels/msa_kv_update.py | 89 + .../kernels/mx_gemv/__init__.py | 4 + .../kernels/mx_gemv/m3_mx_gemv_kernel.py | 1237 ++++ .../kernels/mx_gemv/op.py | 496 ++ .../kernels/triton_sparse_decode.py | 1747 +++++ .../kernels/triton_sparse_decode_mq.py | 237 + .../minimax_m3_nvfp4__sm_107__tp2/modeling.py | 6464 +++++++++++++++++ .../ops/__init__.py | 3 + .../ops/activation/__init__.py | 3 + .../ops/activation/m3_silu_and_mul.py | 83 + .../ops/attention/__init__.py | 3 + .../ops/attention/m3_dense_decode.py | 190 + .../ops/attention/m3_draft_attention.py | 131 + .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 + .../ops/attention/m3_fullctx_decode.py | 142 + .../ops/attention/m3_index_score.py | 239 + .../ops/attention/m3_msa_attention.py | 249 + .../ops/attention/m3_producer.py | 216 + .../ops/attention/m3_qk_norm_rope_fp8.py | 137 + .../ops/attention/m3_qkv_insert.py | 134 + .../ops/attention/m3_select_blocks.py | 82 + .../ops/attention/m3_select_long.py | 192 + .../ops/attention/m3_select_long_dsl.py | 142 + .../ops/attention/m3_sparse_decode.py | 207 + .../ops/attention/m3_sparse_decode_lc.py | 474 ++ .../ops/attention/msa_kv_update.py | 200 + .../ops/comm/__init__.py | 3 + .../ops/comm/m3_allreduce.py | 61 + .../ops/comm/m3_mc_allreduce_norm.py | 319 + .../ops/comm/m3_mc_finalize_allreduce_norm.py | 207 + .../ops/embedding/__init__.py | 3 + .../ops/embedding/m3_embedding.py | 48 + .../ops/gemm/__init__.py | 3 + .../ops/gemm/m3_bf16_linear.py | 41 + .../ops/gemm/m3_dense_swiglu.py | 72 + .../ops/gemm/m3_draft_lowm.py | 173 + .../ops/gemm/m3_fp32_linear.py | 46 + .../ops/gemm/m3_mx_gemv.py | 170 + .../ops/gemm/m3_mx_gemv_bf16.py | 100 + .../ops/gemm/m3_mxfp8_gemm.py | 58 + .../ops/gemm/m3_mxfp8_linear.py | 137 + .../ops/gemm/m3_target_head.py | 346 + .../ops/gemm/m3_ugpu_head.py | 168 + .../ops/moe/__init__.py | 3 + .../ops/moe/m3_moe_engine.py | 342 + .../ops/moe/m3_router_gemv.py | 101 + .../ops/moe/m3_router_gemv_fp4.py | 63 + .../ops/moe/m3_router_tc.py | 114 + .../ops/norm/__init__.py | 3 + .../ops/norm/m3_embedding_firstnorm.py | 98 + .../ops/spec/__init__.py | 4 + .../ops/spec/m3_capture_norm.py | 87 + .../ops/spec/m3_chunk_norm.py | 81 + .../ops/spec/m3_draft_input.py | 140 + .../ops/spec/m3_draft_silu.py | 58 + .../ops/spec/m3_draft_step.py | 432 ++ .../ops/spec/m3_force_vec.py | 136 + .../ops/spec/m3_sampled_accept.py | 308 + .../ops/spec/m3_sampled_accept_ws_words.py | 43 + .../ops/spec/m3_shard_accept.py | 342 + .../ops/spec/m3_slot_scatter.py | 72 + .../ops/spec/m3_spec_offsets.py | 60 + .../ops/spec/m3_step_gather.py | 160 + .../ops/spec/m3_step_prologue.py | 76 + .../minimax_m3_nvfp4__sm_107__tp2/weights.py | 865 +++ .../minimax_m3_nvfp4__sm_107__tp4/__init__.py | 3 + .../kernels/__init__.py | 3 + .../kernels/embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../kernels/embedding_firstnorm/op.py | 142 + .../kernels/long_select.py | 72 + .../kernels/m3_attn_lc.py | 148 + .../kernels/m3_draft_lowm.py | 143 + .../kernels/m3_draft_step.py | 222 + .../kernels/m3_drafter_attn.py | 405 ++ .../kernels/m3_fold_tail.py | 172 + .../kernels/m3_force_vec.py | 64 + .../kernels/m3_fullctx_decode.py | 313 + .../kernels/m3_index_score.py | 638 ++ .../kernels/m3_index_score_t2sc.py | 682 ++ .../kernels/m3_lowm_splitk.py | 1249 ++++ .../kernels/m3_mc_allreduce.py | 422 ++ .../kernels/m3_moe/__init__.py | 16 + .../kernels/m3_moe/m3_moe16_kernel.py | 3871 ++++++++++ .../kernels/m3_moe/m3_moe_kernel.py | 3914 ++++++++++ .../kernels/m3_moe/op.py | 779 ++ .../kernels/m3_moe/op16.py | 288 + .../kernels/m3_prequant.py | 111 + .../kernels/m3_report.py | 34 + .../kernels/m3_sampled_accept.py | 174 + .../kernels/m3_select_long_dsl.py | 367 + .../kernels/m3_shard_accept.py | 176 + .../kernels/m3_spec_tail.py | 116 + .../kernels/m3_target_head.py | 113 + .../kernels/m3_tokmap_guard.py | 185 + .../kernels/msa_kv_update.py | 89 + .../kernels/mx_gemv/__init__.py | 3 + .../kernels/mx_gemv/m3_mx_gemv_kernel.py | 1299 ++++ .../kernels/mx_gemv/op.py | 521 ++ .../kernels/triton_sparse_decode.py | 1730 +++++ .../kernels/triton_sparse_decode_mq.py | 237 + .../minimax_m3_nvfp4__sm_107__tp4/modeling.py | 5941 +++++++++++++++ .../ops/__init__.py | 3 + .../ops/activation/__init__.py | 3 + .../ops/activation/m3_silu_and_mul.py | 83 + .../ops/attention/__init__.py | 3 + .../ops/attention/m3_dense_decode.py | 190 + .../attention/m3_dense_qk_norm_rope_fp8.py | 134 + .../ops/attention/m3_draft_attention.py | 111 + .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 + .../ops/attention/m3_fullctx_decode.py | 142 + .../ops/attention/m3_index_score.py | 159 + .../ops/attention/m3_index_score_t2sc.py | 168 + .../ops/attention/m3_msa_attention.py | 224 + .../ops/attention/m3_producer.py | 223 + .../ops/attention/m3_qk_norm_rope_fp8.py | 137 + .../ops/attention/m3_qkv_insert.py | 134 + .../ops/attention/m3_select_blocks.py | 82 + .../ops/attention/m3_select_long.py | 186 + .../ops/attention/m3_select_long_dsl.py | 140 + .../ops/attention/m3_sparse_decode.py | 215 + .../ops/attention/m3_sparse_decode_lc.py | 436 ++ .../ops/attention/msa_kv_update.py | 200 + .../ops/comm/__init__.py | 3 + .../ops/comm/m3_allreduce.py | 61 + .../ops/comm/m3_mc_allreduce_norm.py | 329 + .../ops/comm/m3_mc_finalize_allreduce_norm.py | 197 + .../ops/embedding/__init__.py | 3 + .../ops/embedding/m3_embedding.py | 48 + .../ops/gemm/__init__.py | 3 + .../ops/gemm/m3_bf16_linear.py | 41 + .../ops/gemm/m3_dense_swiglu.py | 72 + .../ops/gemm/m3_draft_lowm.py | 158 + .../ops/gemm/m3_fp32_linear.py | 46 + .../ops/gemm/m3_mx_gemv.py | 166 + .../ops/gemm/m3_mx_gemv_bf16.py | 106 + .../ops/gemm/m3_mxfp8_gemm.py | 58 + .../ops/gemm/m3_mxfp8_linear.py | 137 + .../ops/gemm/m3_mxfp8_pinned_gemm.py | 69 + .../ops/gemm/m3_mxfp8_pinned_quantize.py | 47 + .../ops/gemm/m3_target_head.py | 346 + .../ops/moe/__init__.py | 3 + .../ops/moe/m3_moe16_engine.py | 67 + .../ops/moe/m3_moe_engine.py | 326 + .../ops/moe/m3_router_tc.py | 119 + .../ops/norm/__init__.py | 3 + .../ops/norm/m3_embedding_firstnorm.py | 87 + .../ops/spec/__init__.py | 4 + .../ops/spec/m3_capture_norm.py | 80 + .../ops/spec/m3_chunk_norm.py | 70 + .../ops/spec/m3_draft_input.py | 140 + .../ops/spec/m3_draft_silu.py | 58 + .../ops/spec/m3_draft_step.py | 417 ++ .../ops/spec/m3_force_vec.py | 86 + .../ops/spec/m3_force_vec_req.py | 62 + .../ops/spec/m3_sampled_accept.py | 289 + .../ops/spec/m3_sampled_accept_ws_words.py | 59 + .../ops/spec/m3_shard_accept.py | 284 + .../ops/spec/m3_slot_scatter.py | 72 + .../ops/spec/m3_spec_offsets.py | 60 + .../ops/spec/m3_step_gather.py | 152 + .../ops/spec/m3_step_prologue.py | 76 + .../minimax_m3_nvfp4__sm_107__tp4/weights.py | 865 +++ .../minimax_m3_nvfp4__sm_107__tp8/__init__.py | 3 + .../kernels/__init__.py | 3 + .../kernels/embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../kernels/embedding_firstnorm/op.py | 142 + .../kernels/long_select.py | 70 + .../kernels/m3_attn_lc8.py | 152 + .../kernels/m3_draft_step.py | 355 + .../kernels/m3_drafter_attn.py | 404 ++ .../kernels/m3_fc_wait.py | 40 + .../kernels/m3_fold_tail.py | 180 + .../kernels/m3_force_vec.py | 63 + .../kernels/m3_fullctx_decode.py | 332 + .../kernels/m3_index_score.py | 692 ++ .../kernels/m3_layer/__init__.py | 14 + .../kernels/m3_layer/phase_producer.py | 487 ++ .../kernels/m3_locality.py | 448 ++ .../kernels/m3_lowm_splitk.py | 1592 ++++ .../kernels/m3_mc_allreduce.py | 439 ++ .../kernels/m3_moe/__init__.py | 16 + .../kernels/m3_moe/m3_moe_kernel.py | 3841 ++++++++++ .../kernels/m3_moe/op.py | 693 ++ .../kernels/m3_prequant.py | 111 + .../kernels/m3_report.py | 34 + .../kernels/m3_router_gemv.py | 107 + .../kernels/m3_sampled_accept.py | 230 + .../kernels/m3_select_long_dsl.py | 367 + .../kernels/m3_shard_accept.py | 176 + .../kernels/m3_spec_tail.py | 113 + .../kernels/m3_target_head.py | 113 + .../kernels/m3_tokmap_guard.py | 185 + .../kernels/msa_kv_update.py | 89 + .../kernels/mx_gemv/__init__.py | 4 + .../kernels/mx_gemv/m3_mx_gemv_kernel.py | 1472 ++++ .../kernels/mx_gemv/op.py | 645 ++ .../kernels/triton_sparse_decode.py | 1754 +++++ .../kernels/triton_sparse_decode_mq.py | 237 + .../minimax_m3_nvfp4__sm_107__tp8/modeling.py | 6151 ++++++++++++++++ .../ops/__init__.py | 3 + .../ops/activation/__init__.py | 3 + .../ops/activation/m3_silu_and_mul.py | 83 + .../ops/attention/__init__.py | 3 + .../ops/attention/m3_dense_decode.py | 190 + .../ops/attention/m3_draft_attention.py | 131 + .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 + .../ops/attention/m3_fullctx_decode.py | 142 + .../ops/attention/m3_index_score.py | 167 + .../ops/attention/m3_msa_attention.py | 249 + .../ops/attention/m3_producer.py | 223 + .../ops/attention/m3_qk_norm_rope_fp8.py | 137 + .../ops/attention/m3_qkv_insert.py | 134 + .../ops/attention/m3_select_blocks.py | 82 + .../ops/attention/m3_select_long.py | 193 + .../ops/attention/m3_select_long_dsl.py | 142 + .../ops/attention/m3_sparse_decode.py | 215 + .../ops/attention/m3_sparse_decode_lc8.py | 487 ++ .../ops/attention/msa_kv_update.py | 200 + .../ops/comm/__init__.py | 3 + .../ops/comm/m3_allreduce.py | 60 + .../ops/comm/m3_mc_allreduce_norm.py | 300 + .../ops/comm/m3_mc_finalize_allreduce_norm.py | 215 + .../ops/embedding/__init__.py | 3 + .../ops/embedding/m3_embedding.py | 48 + .../ops/gemm/__init__.py | 3 + .../ops/gemm/m3_bf16_linear.py | 41 + .../ops/gemm/m3_dense_swiglu.py | 72 + .../ops/gemm/m3_draft_lowm.py | 160 + .../ops/gemm/m3_fp32_linear.py | 46 + .../ops/gemm/m3_mx_gemv.py | 160 + .../ops/gemm/m3_mx_gemv_bf16.py | 100 + .../ops/gemm/m3_mxfp8_gemm.py | 60 + .../ops/gemm/m3_mxfp8_linear.py | 137 + .../ops/gemm/m3_target_head.py | 346 + .../ops/moe/__init__.py | 3 + .../ops/moe/m3_moe_engine.py | 341 + .../ops/moe/m3_router_gemv.py | 101 + .../ops/moe/m3_router_gemv_fp4.py | 63 + .../ops/moe/m3_router_tc.py | 106 + .../ops/norm/__init__.py | 3 + .../ops/norm/m3_embedding_firstnorm.py | 90 + .../ops/spec/__init__.py | 4 + .../ops/spec/m3_capture_norm.py | 80 + .../ops/spec/m3_chunk_norm.py | 74 + .../ops/spec/m3_draft_input.py | 140 + .../ops/spec/m3_draft_silu.py | 58 + .../ops/spec/m3_draft_step.py | 440 ++ .../ops/spec/m3_force_vec.py | 136 + .../ops/spec/m3_sampled_accept.py | 282 + .../ops/spec/m3_sampled_accept_ws_words.py | 59 + .../ops/spec/m3_shard_accept.py | 266 + .../ops/spec/m3_slot_scatter.py | 72 + .../ops/spec/m3_spec_offsets.py | 60 + .../ops/spec/m3_step_gather.py | 160 + .../ops/spec/m3_step_prologue.py | 76 + .../minimax_m3_nvfp4__sm_107__tp8/weights.py | 884 +++ .../modeling_v2/models/minimax_m3/routing.py | 77 +- .../attention/backends/fmha/fallback.py | 13 + .../attention/backends/prims_ts/__init__.py | 2 + .../prims_ts/_block_sparse/compiler.py | 2 + .../backends/prims_ts/_worktile_scope.py | 90 + .../attention/backends/prims_ts/context.py | 3 + .../attention/backends/prims_ts/decode.py | 2 + .../fmha_decode/block_sparse_inspect.py | 3 + .../attention/backends/prims_ts/mla_decode.py | 2 + .../sparse/minimax_m3/cache_manager.py | 117 +- .../sparse/minimax_m3/kernels/paged_cache.py | 39 +- .../kernels/triton_sparse_decode.py | 6 +- .../kernels/trtllm_gen_dense_decode.py | 5 +- .../backends/sparse/minimax_m3/msa_backend.py | 103 +- .../_torch/attention/backends/trtllm.py | 14 +- .../custom_ops/mxfp8_cute_pinned_tactics.py | 160 + .../_torch/custom_ops/torch_custom_ops.py | 35 +- .../cute_dsl_kernels/minimax_m3/__init__.py | 3 + .../embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../minimax_m3/embedding_firstnorm/op.py | 146 + .../minimax_m3/moe/__init__.py | 16 + .../minimax_m3/moe/m3_moe16_kernel.py | 3871 ++++++++++ .../minimax_m3/moe/m3_moe_kernel.py | 3714 ++++++++++ .../cute_dsl_kernels/minimax_m3/moe/op.py | 602 ++ .../cute_dsl_kernels/minimax_m3/moe/op16.py | 288 + .../minimax_m3/mx_gemv/__init__.py | 3 + .../minimax_m3/mx_gemv/m3_mx_gemv_kernel.py | 1297 ++++ .../cute_dsl_kernels/minimax_m3/mx_gemv/op.py | 521 ++ tensorrt_llm/_torch/late_compile.py | 90 + tensorrt_llm/_torch/pyexecutor/_util.py | 13 +- .../_torch/pyexecutor/cuda_graph_runner.py | 18 +- .../engine/runners/decoder/inputs.py | 50 +- .../engine/runners/decoder/metadata.py | 9 + .../engine/runners/decoder/runner.py | 43 +- tensorrt_llm/_torch/pyexecutor/h2d_cache.py | 117 +- .../_torch/pyexecutor/model_loader.py | 10 +- tensorrt_llm/_torch/pyexecutor/profiling.py | 31 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 60 +- .../_torch/pyexecutor/runtime_hooks.py | 62 +- tensorrt_llm/_torch/pyexecutor/step_inputs.py | 302 +- tensorrt_llm/_torch/speculative/interface.py | 17 +- .../_torch/speculative/spec_sampler_base.py | 114 +- .../accuracy/modeling_v2/test_minimax_m3.py | 64 +- .../integration/test_lists/test-db/l0_cpu.yml | 6 + .../msa/test_minimax_m3_kv_size_estimate.py | 84 + .../msa/test_minimax_m3_sparse_attn_decode.py | 10 +- .../attention/sparse/msa/test_msa_backend.py | 77 +- .../sparse/msa/test_msa_runtime_policy.py | 222 +- .../sparse/msa/test_msa_skip_cache_routing.py | 89 + .../test_context_fmha_kernel_presence.py | 31 + .../test_fallback_fmha_extra_page_sizes.py | 58 + .../test_allreduce_tune_nccl_symmetric.py | 34 + .../engine/test_decoder_runtime_hooks.py | 18 +- .../_torch/executor/test_cuda_graph_runner.py | 45 + .../executor/test_first_content_emission.py | 113 + .../_torch/executor/test_h2d_cache.py | 102 + .../_torch/executor/test_h2d_changed_span.py | 157 + .../test_kv_connector_executor_lifetime.py | 1 + .../executor/test_kv_lens_update_once.py | 89 + .../_torch/executor/test_runtime_hooks.py | 9 + .../_torch/executor/test_runtime_opt_ins.py | 58 + .../_torch/executor/test_step_inputs.py | 340 + .../kernels/parallel/test_mxfp8_mxfp8_gemm.py | 17 +- .../test_modeling_v2_m3_silu_and_mul.py | 49 + .../attention/_m3_attn_runtime_tp2.py | 275 + .../attention/_m3_nvfp4_cache.py | 108 + ...est_modeling_v2_m3_block_table_contract.py | 106 + ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 143 + .../test_modeling_v2_m3_draft_attention.py | 148 + ...ing_v2_m3_draft_attention_paged_context.py | 163 + .../test_modeling_v2_m3_drafter_fullctx.py | 271 + ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 406 ++ ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 49 + .../test_modeling_v2_m3_fp8_saturate.py | 102 + .../test_modeling_v2_m3_fullctx_decode.py | 215 + .../test_modeling_v2_m3_identity_selection.py | 102 + .../test_modeling_v2_m3_index_score.py | 107 + .../test_modeling_v2_m3_msa_attention.py | 224 + .../test_modeling_v2_m3_n_valid_contract.py | 138 + .../attention/test_modeling_v2_m3_producer.py | 218 + .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 81 + .../test_modeling_v2_m3_qkv_insert.py | 80 + .../test_modeling_v2_m3_select_blocks.py | 45 + .../test_modeling_v2_m3_select_long.py | 179 + .../test_modeling_v2_m3_select_long_dsl.py | 191 + .../test_modeling_v2_m3_sparse_decode.py | 451 ++ .../test_modeling_v2_m3_sparse_decode_lc.py | 490 ++ .../test_modeling_v2_m3_tokmap_guard.py | 226 + .../test_modeling_v2_msa_kv_update.py | 128 + .../comm/_m3_allreduce_op_matrix.py | 440 ++ .../comm/_m3_mc_allreduce_norm_op_matrix.py | 755 ++ .../comm/_m3_moe_push_op_matrix.py | 411 ++ .../comm/_rank_job.py | 87 + .../comm/test_modeling_v2_m3_allreduce.py | 39 + ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 21 + ...deling_v2_m3_mc_finalize_allreduce_norm.py | 33 + .../test_modeling_v2_m3_moe_push_op_matrix.py | 25 + .../test_modeling_v2_m3_embedding.py | 37 + .../gemm/_m3_mxfp8_operands.py | 30 + .../gemm/test_modeling_v2_m3_bf16_linear.py | 31 + .../gemm/test_modeling_v2_m3_dense_seam.py | 142 + .../gemm/test_modeling_v2_m3_dense_swiglu.py | 142 + .../gemm/test_modeling_v2_m3_draft_lowm.py | 183 + .../gemm/test_modeling_v2_m3_fp32_linear.py | 32 + .../gemm/test_modeling_v2_m3_mx_gemv.py | 87 + .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 63 + .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 38 + .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 46 + .../gemm/test_modeling_v2_m3_pinned_gemm.py | 49 + .../gemm/test_modeling_v2_m3_ugpu_head.py | 116 + .../gemm/test_r1b1_gemv_order_gate_a.py | 82 + .../moe/_m3_moe_operands.py | 209 + .../moe/_m3_router.py | 86 + .../test_modeling_v2_m3_fp32_router_rows.py | 36 + .../moe/test_modeling_v2_m3_moe_engine.py | 170 + .../moe/test_modeling_v2_m3_moe_epoch_wrap.py | 59 + .../test_modeling_v2_m3_moe_finalize_rows.py | 63 + .../moe/test_modeling_v2_m3_moe_flat_gate.py | 97 + .../moe/test_modeling_v2_m3_moe_sh_take.py | 73 + ...test_modeling_v2_m3_moe_triton_finalize.py | 74 + .../moe/test_modeling_v2_m3_router_tc.py | 58 + ...test_modeling_v2_m3_embedding_firstnorm.py | 103 + .../spec/_m3_spec_inputs.py | 106 + .../spec/test_modeling_v2_m3_capture_norm.py | 57 + .../test_modeling_v2_m3_capture_workspaces.py | 56 + .../spec/test_modeling_v2_m3_chunk_norm.py | 50 + ...test_modeling_v2_m3_draft_ar_push_words.py | 107 + .../spec/test_modeling_v2_m3_draft_input.py | 98 + .../test_modeling_v2_m3_draft_push_take.py | 43 + .../spec/test_modeling_v2_m3_draft_silu.py | 88 + .../spec/test_modeling_v2_m3_draft_step.py | 814 +++ .../spec/test_modeling_v2_m3_fold_tail.py | 295 + .../spec/test_modeling_v2_m3_force_al.py | 63 + .../spec/test_modeling_v2_m3_force_vec.py | 47 + .../spec/test_modeling_v2_m3_force_vec_req.py | 48 + .../spec/test_modeling_v2_m3_lowm_splitk.py | 65 + .../test_modeling_v2_m3_sampled_accept.py | 276 + .../spec/test_modeling_v2_m3_shard_accept.py | 127 + .../spec/test_modeling_v2_m3_slot_scatter.py | 57 + .../spec/test_modeling_v2_m3_spec_offsets.py | 138 + .../spec/test_modeling_v2_m3_step_gather.py | 80 + .../spec/test_modeling_v2_m3_step_prologue.py | 192 + .../test_modeling_v2_m3_wide_draft_logits.py | 132 + .../test_modeling_v2_m3_capture_boundary.py | 163 + .../test_modeling_v2_m3_decode_dense_route.py | 85 + .../test_modeling_v2_m3_dispatch.py | 771 ++ ...odeling_v2_m3_draft_attention_bcg_route.py | 69 + ...st_modeling_v2_m3_drafter_fullctx_cells.py | 34 + .../test_modeling_v2_m3_rank_agreement.py | 179 + .../test_modeling_v2_m3_wide_lc_cap.py | 46 + .../test_modeling_v2_m3_silu_and_mul.py | 49 + .../attention/_m3_attn_runtime_tp4.py | 274 + .../attention/_m3_nvfp4_cache.py | 104 + .../test_modeling_v2_m3_dense_decode.py | 153 + ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 180 + .../test_modeling_v2_m3_draft_attention.py | 138 + .../test_modeling_v2_m3_drafter_fullctx.py | 289 + ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 371 + ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 49 + .../test_modeling_v2_m3_fp8_saturate.py | 102 + .../test_modeling_v2_m3_fullctx_decode.py | 312 + .../test_modeling_v2_m3_index_score.py | 133 + .../test_modeling_v2_m3_index_score_t2sc.py | 112 + .../test_modeling_v2_m3_msa_attention.py | 214 + .../test_modeling_v2_m3_n_valid_contract.py | 138 + .../attention/test_modeling_v2_m3_producer.py | 209 + .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 81 + .../test_modeling_v2_m3_qkv_insert.py | 80 + .../test_modeling_v2_m3_select_blocks.py | 45 + .../test_modeling_v2_m3_select_long.py | 183 + .../test_modeling_v2_m3_select_long_dsl.py | 191 + .../test_modeling_v2_m3_sparse_decode.py | 442 ++ .../test_modeling_v2_m3_sparse_decode_lc.py | 526 ++ .../test_modeling_v2_m3_tokmap_guard.py | 226 + .../test_modeling_v2_msa_kv_update.py | 133 + .../comm/_m3_allreduce_op_matrix.py | 440 ++ .../comm/_m3_capture_in_boundary_op_matrix.py | 307 + .../comm/_m3_draft_ar_norm_op_matrix.py | 533 ++ .../comm/_m3_mc_allreduce_norm_op_matrix.py | 665 ++ .../comm/_rank_job.py | 87 + .../comm/test_modeling_v2_m3_allreduce.py | 39 + ...test_modeling_v2_m3_capture_in_boundary.py | 22 + .../comm/test_modeling_v2_m3_draft_ar_norm.py | 32 + ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 20 + ...deling_v2_m3_mc_finalize_allreduce_norm.py | 48 + .../test_modeling_v2_m3_embedding.py | 37 + .../gemm/_m3_mxfp8_operands.py | 30 + .../gemm/test_modeling_v2_m3_bf16_linear.py | 31 + .../gemm/test_modeling_v2_m3_dense_swiglu.py | 142 + .../gemm/test_modeling_v2_m3_draft_lowm.py | 182 + .../gemm/test_modeling_v2_m3_fp32_linear.py | 32 + .../gemm/test_modeling_v2_m3_mx_gemv.py | 176 + .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 68 + .../gemm/test_modeling_v2_m3_mx_gemv_order.py | 78 + .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 38 + .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 46 + .../test_modeling_v2_m3_mxfp8_pinned_gemm.py | 82 + ...st_modeling_v2_m3_mxfp8_pinned_quantize.py | 48 + .../gemm/test_modeling_v2_m3_target_head.py | 253 + .../moe/_m3_moe_operands.py | 170 + .../moe/_m3_router.py | 39 + .../moe/test_modeling_v2_m3_moe16_engine.py | 63 + .../moe/test_modeling_v2_m3_moe_engine.py | 139 + .../test_modeling_v2_m3_moe_flat_gate_tp4.py | 74 + .../moe/test_modeling_v2_m3_router_tc.py | 60 + ...test_modeling_v2_m3_embedding_firstnorm.py | 76 + .../spec/_m3_spec_inputs.py | 106 + .../spec/test_modeling_v2_m3_capture_norm.py | 62 + .../test_modeling_v2_m3_capture_normed.py | 70 + .../spec/test_modeling_v2_m3_chunk_norm.py | 55 + .../spec/test_modeling_v2_m3_draft_input.py | 103 + .../test_modeling_v2_m3_draft_push_take.py | 42 + .../spec/test_modeling_v2_m3_draft_silu.py | 98 + .../spec/test_modeling_v2_m3_draft_step.py | 694 ++ ...t_modeling_v2_m3_draft_step_fixed_order.py | 148 + .../test_modeling_v2_m3_exact_numerics.py | 289 + .../spec/test_modeling_v2_m3_fold_tail.py | 302 + .../spec/test_modeling_v2_m3_force_vec.py | 47 + .../spec/test_modeling_v2_m3_force_vec_req.py | 93 + .../spec/test_modeling_v2_m3_lowm_splitk.py | 63 + .../test_modeling_v2_m3_sampled_accept.py | 297 + .../test_modeling_v2_m3_sampled_accept_1r.py | 186 + .../spec/test_modeling_v2_m3_shard_accept.py | 125 + .../spec/test_modeling_v2_m3_slot_scatter.py | 62 + .../spec/test_modeling_v2_m3_spec_offsets.py | 138 + .../spec/test_modeling_v2_m3_step_gather.py | 85 + .../spec/test_modeling_v2_m3_step_prologue.py | 192 + .../test_modeling_v2_m3_capture_boundary.py | 154 + .../test_modeling_v2_m3_dispatch.py | 286 + .../test_modeling_v2_m3_prebuild.py | 314 + .../test_modeling_v2_m3_rank_agreement.py | 181 + .../test_modeling_v2_m3_silu_and_mul.py | 49 + .../attention/_m3_attn_runtime_tp8.py | 274 + .../attention/_m3_nvfp4_cache.py | 105 + ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 144 + .../test_modeling_v2_m3_draft_attention.py | 148 + .../test_modeling_v2_m3_drafter_fullctx.py | 271 + ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 371 + ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 49 + .../test_modeling_v2_m3_fp8_saturate.py | 102 + .../test_modeling_v2_m3_fullctx_decode.py | 310 + .../test_modeling_v2_m3_identity_selection.py | 103 + .../test_modeling_v2_m3_index_score.py | 111 + .../test_modeling_v2_m3_msa_attention.py | 224 + .../test_modeling_v2_m3_n_valid_contract.py | 138 + .../attention/test_modeling_v2_m3_producer.py | 209 + .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 81 + .../test_modeling_v2_m3_qkv_insert.py | 80 + .../test_modeling_v2_m3_select_blocks.py | 45 + .../test_modeling_v2_m3_select_long.py | 179 + .../test_modeling_v2_m3_select_long_dsl.py | 191 + .../test_modeling_v2_m3_sparse_decode.py | 626 ++ .../test_modeling_v2_m3_sparse_decode_lc8.py | 360 + .../test_modeling_v2_m3_tokmap_guard.py | 226 + .../test_modeling_v2_msa_kv_update.py | 128 + .../comm/_m3_allreduce_op_matrix.py | 454 ++ .../comm/_m3_mc_allreduce_norm_op_matrix.py | 788 ++ .../comm/_m3_moe_push_op_matrix.py | 403 + .../comm/_rank_job.py | 111 + .../comm/test_modeling_v2_m3_allreduce.py | 42 + ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 62 + .../test_modeling_v2_m3_embedding.py | 37 + .../gemm/_m3_mxfp8_operands.py | 30 + .../gemm/test_modeling_v2_m3_bf16_linear.py | 31 + .../gemm/test_modeling_v2_m3_dense_swiglu.py | 142 + .../gemm/test_modeling_v2_m3_draft_lowm.py | 177 + .../gemm/test_modeling_v2_m3_fp32_linear.py | 32 + .../gemm/test_modeling_v2_m3_mx_gemv.py | 87 + .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 63 + .../test_modeling_v2_m3_mx_gemv_producer.py | 166 + .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 38 + .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 46 + .../gemm/test_modeling_v2_m3_pinned_gemm.py | 52 + .../test_modeling_v2_m3_swiglu_oai_bf16x2.py | 149 + .../gemm/test_r1b1_gemv_order_gate_a.py | 85 + .../moe/_m3_moe_operands.py | 172 + .../moe/_m3_router.py | 86 + .../moe/test_modeling_v2_m3_moe_engine.py | 139 + .../moe/test_modeling_v2_m3_router_tc.py | 58 + ...test_modeling_v2_m3_embedding_firstnorm.py | 102 + .../spec/_m3_spec_inputs.py | 106 + .../spec/test_modeling_v2_m3_capture_norm.py | 57 + .../test_modeling_v2_m3_capture_normed.py | 71 + .../test_modeling_v2_m3_capture_target.py | 82 + .../test_modeling_v2_m3_capture_workspaces.py | 56 + .../spec/test_modeling_v2_m3_chunk_norm.py | 50 + .../spec/test_modeling_v2_m3_di_fc_flag.py | 91 + .../spec/test_modeling_v2_m3_draft_input.py | 98 + .../test_modeling_v2_m3_draft_push_take.py | 43 + .../spec/test_modeling_v2_m3_draft_silu.py | 88 + .../spec/test_modeling_v2_m3_draft_step.py | 747 ++ ...t_modeling_v2_m3_draft_step_fixed_order.py | 150 + .../spec/test_modeling_v2_m3_fold_tail.py | 376 + .../spec/test_modeling_v2_m3_force_al.py | 63 + .../spec/test_modeling_v2_m3_force_vec.py | 47 + .../spec/test_modeling_v2_m3_force_vec_req.py | 48 + .../spec/test_modeling_v2_m3_lowm_splitk.py | 98 + .../test_modeling_v2_m3_sampled_accept.py | 267 + .../test_modeling_v2_m3_sampled_accept_1r.py | 186 + .../spec/test_modeling_v2_m3_shard_accept.py | 120 + .../spec/test_modeling_v2_m3_slot_scatter.py | 57 + .../spec/test_modeling_v2_m3_spec_offsets.py | 138 + .../spec/test_modeling_v2_m3_step_gather.py | 80 + .../spec/test_modeling_v2_m3_step_prologue.py | 228 + .../test_modeling_v2_m3_capture_boundary.py | 103 + .../test_modeling_v2_m3_dispatch.py | 479 ++ .../test_modeling_v2_m3_rank_agreement.py | 190 + .../modeling_v2/test_modeling_v2_routing.py | 72 +- .../test_one_model_sampling_rejection.py | 29 + .../hw_agnostic/test_ragged_spec_sampler.py | 2 +- .../speculative/test_rng_upload_pack.py | 93 + .../test_spec_sampler_step_fold.py | 116 + tests/unittest/_torch/test_model_config.py | 27 +- .../test_mxfp8_cute_pinned_tactics.py | 100 + 737 files changed, 168821 insertions(+), 205 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DraftStepOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SelectLongOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DraftStepOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4McAllreduceNormOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SelectLongOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/long_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_attn_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_drafter_attn.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fold_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_locality.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_direct.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_splitk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_mc_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe_finalize.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_report.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_spec_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_tokmap_guard.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/m3_mx_gemv_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept_ws_words.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/embedding_firstnorm/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/long_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_attn_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_drafter_attn.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fold_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score_t2sc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_lowm_splitk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_mc_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe16_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_report.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_spec_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_tokmap_guard.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_qk_norm_rope_fp8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept_ws_words.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp8/kernels/embedding_firstnorm/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/long_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_attn_lc8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_drafter_attn.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fc_wait.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fold_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/phase_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_locality.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_lowm_splitk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_mc_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_report.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_spec_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_tokmap_guard.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp8/kernels/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept_ws_words.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/weights.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/_worktile_scope.py create mode 100644 tensorrt_llm/_torch/custom_ops/mxfp8_cute_pinned_tactics.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/late_compile.py create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_msa_skip_cache_routing.py create mode 100644 tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py create mode 100644 tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py create mode 100644 tests/unittest/_torch/executor/test_first_content_emission.py create mode 100644 tests/unittest/_torch/executor/test_h2d_changed_span.py create mode 100644 tests/unittest/_torch/executor/test_kv_lens_update_once.py create mode 100644 tests/unittest/_torch/executor/test_runtime_opt_ins.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_attn_runtime_tp2.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_nvfp4_cache.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_identity_selection.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_allreduce.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_moe_push_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_moe_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_push_take.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_attn_runtime_tp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_nvfp4_cache.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_saturate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_allreduce.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_order.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_moe_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_normed.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept_1r.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_attn_runtime_tp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_nvfp4_cache.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_saturate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_identity_selection.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_allreduce.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_moe_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_push_take.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py create mode 100644 tests/unittest/_torch/speculative/test_rng_upload_pack.py create mode 100644 tests/unittest/_torch/thop/parallel/test_mxfp8_cute_pinned_tactics.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 57ae33cb1935..936fdba06042 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,3 +1,132 @@ +diff --git a/__init__.py b/__init__.py +index e40d8ce0d41c48ad17099e097c64a70e1f6d7a4b..c49ad2fc908d8794105d7340d7614347db123c53 100644 +--- a/__init__.py ++++ b/__init__.py +@@ -14,6 +14,8 @@ + + """Experimental task-scheduled attention entry points.""" + ++from . import _worktile_scope # noqa: F401 (first: scopes task scheduling's WorkTileInfo rebind) ++ + from .block_sparse import ( + BlockSparsePagedTSWrapper, + BlockSparseTSWrapper, +diff --git a/_block_sparse/compiler.py b/_block_sparse/compiler.py +index 1a9d42857f94459f8d2bb8bb523080c0b6bf1ea7..d557504e26512148da732a891692e9c5ea295b90 100644 +--- a/_block_sparse/compiler.py ++++ b/_block_sparse/compiler.py +@@ -14,6 +14,7 @@ + + """Compilation and caching for PrimTS block-sparse attention adapters.""" + ++from .._worktile_scope import with_scalar_work_tile_info + from collections.abc import Callable + import functools + +@@ -29,6 +30,7 @@ from .config import _BlockSparseCompileKey, _make_block_sparse_config + _COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 3" + + ++@with_scalar_work_tile_info + def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: + """Compile one contiguous or paged attention adapter for ``key``. + +diff --git a/_worktile_scope.py b/_worktile_scope.py +new file mode 100644 +index 0000000000000000000000000000000000000000..a54e0d6f9893a27db336dda5852c26aa75055b21 +--- /dev/null ++++ b/_worktile_scope.py +@@ -0,0 +1,90 @@ ++# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ++# SPDX-License-Identifier: Apache-2.0 ++"""Keep CuTe DSL task scheduling's WorkTileInfo rebind inside PrimTS. ++ ++Importing ``cutlass.experimental.task_scheduling`` rebinds six attributes of CUTLASS' ++``WorkTileInfo`` (``__init__``, the two MLIR value hooks, ``tile_idx``, ``is_valid_tile`` and ++``update_from``) to a scalar three-axis form, for every kernel in the process. Kernels that ++subclass ``WorkTileInfo`` with a different tile shape (FA4's four-axis tile) then fail to trace. ++ ++This module is imported first by the package. It snapshots the original attributes, imports ++the task scheduling resources (which apply the rebind), snapshots the scalar set, and puts the ++originals back. ``scalar_work_tile_info`` installs the scalar set only while a PrimTS kernel ++traces and compiles; compiled kernels do not call these Python methods. Nothing about the ++generated code changes. Kill switch: TRTLLM_PRIMS_TS_SCOPED_WORKTILE=0 keeps the global rebind. ++""" ++ ++import contextlib ++import functools ++import os ++import sys ++import threading ++ ++from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo as _WorkTileInfo ++ ++_ATTRS = ( ++ "__init__", ++ "__extract_mlir_values__", ++ "__new_from_mlir_values__", ++ "tile_idx", ++ "is_valid_tile", ++ "update_from", ++) ++_MISSING = object() ++_RESOURCES = "cutlass.experimental.task_scheduling.resources" ++ ++ ++def _snapshot() -> dict: ++ return {name: _WorkTileInfo.__dict__.get(name, _MISSING) for name in _ATTRS} ++ ++ ++def _install(attrs: dict) -> None: ++ for name, value in attrs.items(): ++ if value is _MISSING: ++ if name in _WorkTileInfo.__dict__: ++ delattr(_WorkTileInfo, name) ++ else: ++ setattr(_WorkTileInfo, name, value) ++ ++ ++# Only scope the rebind when this module saw the originals, i.e. task scheduling was not imported ++# before PrimTS. Otherwise keep whatever is installed (the previous global behavior). ++_ORIGINAL = None if _RESOURCES in sys.modules else _snapshot() ++import cutlass.experimental.task_scheduling.resources # noqa: E402,F401 (applies the scalar rebind) ++ ++_SCALAR = _snapshot() ++SCOPED = _ORIGINAL is not None and os.environ.get("TRTLLM_PRIMS_TS_SCOPED_WORKTILE", "1") != "0" ++_LOCK = threading.RLock() ++_DEPTH = 0 ++if SCOPED: ++ _install(_ORIGINAL) ++ ++ ++@contextlib.contextmanager ++def scalar_work_tile_info(): ++ """Install the task scheduling WorkTileInfo form while a PrimTS kernel traces (re-entrant).""" ++ global _DEPTH ++ if not SCOPED: ++ yield ++ return ++ with _LOCK: ++ if _DEPTH == 0: ++ _install(_SCALAR) ++ _DEPTH += 1 ++ try: ++ yield ++ finally: ++ _DEPTH -= 1 ++ if _DEPTH == 0: ++ _install(_ORIGINAL) ++ ++ ++def with_scalar_work_tile_info(fn): ++ """Decorator form of ``scalar_work_tile_info`` for the PrimTS compile functions.""" ++ ++ @functools.wraps(fn) ++ def wrapper(*args, **kwargs): ++ with scalar_work_tile_info(): ++ return fn(*args, **kwargs) ++ ++ return wrapper diff --git a/block_sparse.py b/block_sparse.py index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea063f9e6e4 100644 --- a/block_sparse.py @@ -61,10 +190,18 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415450369df 100644 +index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..171c02224d12fd7d0481a7ba1d905fc8281b1ec4 100644 --- a/context.py +++ b/context.py -@@ -46,7 +46,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional +@@ -35,6 +35,7 @@ now; unlike the decode APIs, their ``@flashinfer_experimental_api`` decorators d + register trace templates. + """ + ++from ._worktile_scope import with_scalar_work_tile_info + from dataclasses import dataclass + import functools + import importlib.metadata +@@ -46,7 +47,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional import torch @@ -73,7 +210,7 @@ index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415 if TYPE_CHECKING: -@@ -488,7 +488,7 @@ def _validate_device(device: torch.device) -> int: +@@ -488,7 +489,7 @@ def _validate_device(device: torch.device) -> int: # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. if capability == (10, 7): @@ -82,8 +219,24 @@ index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415 require_cute_dsl_arch(device_index) return device_index +@@ -1753,6 +1754,7 @@ def _paged_context_compile_spec( + + + @functools.cache ++@with_scalar_work_tile_info + def _get_compiled_context( + compile_spec: _ContextCompileSpec, + ): +@@ -1973,6 +1975,7 @@ def _get_compiled_context( + + + @functools.cache ++@with_scalar_work_tile_info + def _get_compiled_paged_context( + compile_spec: _PagedContextCompileSpec, + ): diff --git a/decode.py b/decode.py -index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa6d76a609 100644 +index 41479a386d72c27eb0df76107318dc09256aff06..a161190007f4a661035d3e555db2d4a497e5228e 100644 --- a/decode.py +++ b/decode.py @@ -1,3 +1,4 @@ @@ -91,7 +244,15 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -24,11 +25,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast +@@ -14,6 +15,7 @@ + + """Task-scheduled paged decode with a FlashInfer-style plan/run lifecycle.""" + ++from ._worktile_scope import with_scalar_work_tile_info + from collections.abc import Callable, Sequence + from dataclasses import dataclass + import functools +@@ -24,11 +26,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast import torch @@ -104,7 +265,7 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa from . import _q_token_kv_block_sparse_policy as _sparse_policy from ._block_sparse.common import _num_sparse_pattern_heads -@@ -789,7 +786,7 @@ def _validate_runtime_device(device: torch.device) -> None: +@@ -789,7 +787,7 @@ def _validate_runtime_device(device: torch.device) -> None: # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. if capability == (10, 7): @@ -113,7 +274,15 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa require_cute_dsl_arch(device_index) -@@ -4046,7 +4043,7 @@ class BatchDecodePagedTSWrapper: +@@ -1513,6 +1511,7 @@ def _make_decode_compile_spec( + + + @functools.cache ++@with_scalar_work_tile_info + def _get_compiled_decode( + compile_spec: _DecodeCompileSpec, + ): +@@ -4046,7 +4045,7 @@ class BatchDecodePagedTSWrapper: # previous complete plan revision usable. self._plan_state = candidate @@ -122,7 +291,7 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa def run( self, q: torch.Tensor, -@@ -4228,7 +4225,7 @@ class BatchDecodePagedTSWrapper: +@@ -4228,7 +4227,7 @@ class BatchDecodePagedTSWrapper: ) @@ -200,8 +369,36 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb + _load_runtime_seq_len_kv, + _resolve_kv_tile_idx_context, +) +diff --git a/kernels/fmha_decode/block_sparse_inspect.py b/kernels/fmha_decode/block_sparse_inspect.py +index 5728a7c6d373ddd694fa0c740dd01f672a65f648..e5f8de86571f54af313adcc0493bc65e4c9ce055 100644 +--- a/kernels/fmha_decode/block_sparse_inspect.py ++++ b/kernels/fmha_decode/block_sparse_inspect.py +@@ -26,6 +26,7 @@ Both publish one validation status plus the maximum row width in one Int64 + summary; no route payload is constructed. + """ + ++from ..._worktile_scope import with_scalar_work_tile_info + import functools + from collections.abc import Callable + +@@ -450,6 +451,7 @@ def _fake_compact( + + + @functools.cache ++@with_scalar_work_tile_info + def compile_block_sparse_inspection( + *, + device_index: int, +@@ -495,6 +497,7 @@ def compile_block_sparse_inspection( + + + @functools.cache ++@with_scalar_work_tile_info + def compile_paged_block_sparse_metadata_inspection( + *, + device_index: int, diff --git a/mla_decode.py b/mla_decode.py -index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2eef74a1a4b 100644 +index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..0e82e5330728f78aa93709145382be5f2a015feb 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -209,7 +406,15 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2ee # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -21,10 +22,9 @@ from typing import Any, Literal, Optional, cast +@@ -14,6 +15,7 @@ + + """Task-scheduled paged MLA decode with a plan/run lifecycle.""" + ++from ._worktile_scope import with_scalar_work_tile_info + from collections.abc import Callable + from dataclasses import dataclass, field + import functools +@@ -21,10 +23,9 @@ from typing import Any, Literal, Optional, cast import torch @@ -221,7 +426,15 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2ee ) from .decode import ( -@@ -1576,7 +1576,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1044,6 +1045,7 @@ def _make_mla_decode_compile_spec( + + + @functools.cache ++@with_scalar_work_tile_info + def _get_compiled_mla_decode( + compile_spec: _MLADecodeCompileSpec, + ): +@@ -1576,7 +1578,7 @@ class BatchMLADecodePagedTSWrapper: split_kv=int(dict(policy)["split_kv"]), ) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 617b168021c6..375e9259006d 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -9,8 +9,8 @@ vendors: include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch - patch_digest: sha256:a4360ca506d25545db408d22875e678871970fec94bdc1166dd6108d3b504d1c - digest: sha256-tree-v1:1bd57713d2487e94c0fc4b134bed4f7eff275282cb8ef471775e87db503ded03 + patch_digest: sha256:e098557cdac388da6b8b1a5c055d5c564a354273c083563ef5f581fb8df5e37b + digest: sha256-tree-v1:c30bd5b6865fe65149c283088ab8b9b3c4c8c544495793e4e224e91d45d632df openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 9cbf1480a879..7245364e01e6 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -81,6 +81,9 @@ list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") # architecture-scoped OBJECT library below. list(FILTER SRC_CU EXCLUDE REGEX "minimaxM3/(fusedDecodeAttn|attnMega|attnC1)\\.cu$") +# The MiniMax-M3 sm_107 decode kernels are built as their own sm_107-only OBJECT +# library below. +list(FILTER SRC_CU EXCLUDE REGEX "minimaxM3Sm107Tp[248]/.*") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") @@ -126,11 +129,38 @@ set_property(TARGET minimax_m3_attn_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET minimax_m3_attn_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) set_cuda_architectures(minimax_m3_attn_src 100 103) +# MiniMax-M3 sm_107 decode kernels (minimaxM3Sm107Tp{2,4,8}/): written for +# sm_107 only (PDL griddepcontrol, sm_10x-only PTX), so they are compiled for +# the sm_10x entries of this build's architectures only (the same targets, e.g. +# the sm_100f family, that kernels_src compiles its sm_10x code for), not for +# the older ones. A build without an sm_10x architecture compiles them for +# sm_107a: their torch ops link against the host launchers, and only the sm_107 +# targets launch them. +file(GLOB MINIMAX_M3_SM107_SRC minimaxM3Sm107Tp2/*.cu minimaxM3Sm107Tp4/*.cu + minimaxM3Sm107Tp8/*.cu) +add_library(minimax_m3_sm107_src OBJECT ${MINIMAX_M3_SM107_SRC}) +set_property(TARGET minimax_m3_sm107_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET minimax_m3_sm107_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS + ON) +target_include_directories( + minimax_m3_sm107_src + PRIVATE + $ +) +target_link_libraries(minimax_m3_sm107_src PRIVATE trtllm_gen_fmha_interface) +set(MINIMAX_M3_SM107_ARCHS ${CMAKE_CUDA_ARCHITECTURES}) +list(FILTER MINIMAX_M3_SM107_ARCHS INCLUDE REGEX "^10[0-9]") +if(NOT MINIMAX_M3_SM107_ARCHS) + set(MINIMAX_M3_SM107_ARCHS 107a-real) +endif() +set_property(TARGET minimax_m3_sm107_src PROPERTY CUDA_ARCHITECTURES + ${MINIMAX_M3_SM107_ARCHS}) + add_library( kernels_src STATIC ${SRC_CPP} ${SRC_CU} ${MOE_KERNELS_SRC} $<$:$> - $) + $ $) set_property(TARGET kernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) target_include_directories( diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu index d12b765e5ecf..9ef178897f9f 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu @@ -520,36 +520,13 @@ void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, i launchGemv<1>(x, ldx, m, k, w, e, out, qOut, sfOut, sfScale, pdl, stream); } -void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out, - uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream) +int minimaxM3RouterTcCluster(int32_t k) { - int64_t const M = m, K = k; - TLLM_CHECK_WITH_INFO( - M >= 1 && M <= kTcMaxM && K % (kTcSlices * 64) == 0, "unsupported router shape M=%d K=%d", m, k); - int const ks = K / kTcSlices; - auto launch_tc = [&](auto kernel, int cl) -> cudaError_t - { - int const kc = (kTcSlices / cl) * ks; - size_t const smem = static_cast(kTcExperts) * (kc + 4) * 4 + static_cast(M) * (kc + 8) * 2; - cudaLaunchConfig_t cfg = {}; - cfg.gridDim = dim3(8 * cl); - cfg.blockDim = dim3(256); - cfg.dynamicSmemBytes = smem; - cfg.stream = stream; - cudaLaunchAttribute attr[2]; - attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - attr[1].id = cudaLaunchAttributeClusterDimension; - attr[1].val.clusterDim.x = cl; - attr[1].val.clusterDim.y = 1; - attr[1].val.clusterDim.z = 1; - cfg.attrs = attr; - cfg.numAttrs = 2; - return cudaLaunchKernelEx( - &cfg, kernel, x, ldx, w, out, static_cast(M), static_cast(K), qOut, sfOut, sfScale); - }; - // 16-CTA clusters (128 CTAs, one K slice each) when the device co-schedules all 8, else 8-CTA clusters. - static int const cluster = [&]() + // 16-CTA clusters (128 CTAs, one K slice each) when the device co-schedules all 8, else 8-CTA clusters; decided + // once per process (device-static, for the first K seen: MiniMax-M3 has one hidden size) and configures both + // kernels. + int const ks = k / kTcSlices; + static int const cluster = [ks]() { int device = 0, max_optin = 0; cudaGetDevice(&device); @@ -592,6 +569,38 @@ void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int } return clusters >= 8 ? 16 : 8; }(); + return cluster; +} + +void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out, + uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream) +{ + int64_t const M = m, K = k; + TLLM_CHECK_WITH_INFO( + M >= 1 && M <= kTcMaxM && K % (kTcSlices * 64) == 0, "unsupported router shape M=%d K=%d", m, k); + int const ks = K / kTcSlices; + auto launch_tc = [&](auto kernel, int cl) -> cudaError_t + { + int const kc = (kTcSlices / cl) * ks; + size_t const smem = static_cast(kTcExperts) * (kc + 4) * 4 + static_cast(M) * (kc + 8) * 2; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(8 * cl); + cfg.blockDim = dim3(256); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = cl; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + return cudaLaunchKernelEx( + &cfg, kernel, x, ldx, w, out, static_cast(M), static_cast(K), qOut, sfOut, sfScale); + }; + int const cluster = minimaxM3RouterTcCluster(k); cudaError_t const err = cluster == 16 ? launch_tc(m3_router_tc_kernel<16>, 16) : launch_tc(m3_router_tc_kernel<8>, 8); if (err != cudaSuccess) diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h index 6f4467883b24..6548cb8ca83d 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h @@ -41,6 +41,10 @@ void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, i void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out, uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream); +//! The cluster size `invokeMinimaxM3RouterTc` launches with (16 or 8): decided once per process from the device's +//! co-scheduling (device-static) and configures both kernels. SM107 callers use the router only when this is 16. +int minimaxM3RouterTcCluster(int32_t k); + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu index 4069da14faf0..45526249947f 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu @@ -43,6 +43,9 @@ namespace specTail namespace { constexpr int kVec = 8; // bf16 per 16-byte access +// draft_input loads its constant norm weight row before the grid dependency wait (the same values, bit for bit). +// Kill: false (load the row after the wait, with the normalization). +constexpr bool kDraftInputPrewaitWeight = true; __device__ __forceinline__ float block_sum(float v, float* red) { @@ -72,6 +75,20 @@ __global__ void __launch_bounds__(kThreads) __shared__ float red[kThreads / 32]; int const row = blockIdx.x; bool const is_hidden = blockIdx.y == 1; + // The norm weight row is constant (no grid of the step writes it): load it before the grid dependency wait, so + // after the wait only the dependent loads (the id / hidden row) remain ahead of the reduction. The same values and + // arithmetic. + __nv_bfloat16 const* const w = is_hidden ? w_hidden : w_embed; + uint4 wv[kIters]; + if constexpr (kDraftInputPrewaitWeight) + { +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const c = (it * kThreads + threadIdx.x) * kVec; + wv[it] = c < H ? *reinterpret_cast(w + c) : make_uint4(0u, 0u, 0u, 0u); + } + } #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("griddepcontrol.wait;" ::: "memory"); #endif @@ -112,7 +129,6 @@ __global__ void __launch_bounds__(kThreads) asm volatile("griddepcontrol.launch_dependents;"); #endif float const rs = rsqrtf(block_sum(ss, red) / H + (is_hidden ? eps_hidden : eps_embed)); - __nv_bfloat16 const* w = is_hidden ? w_hidden : w_embed; __nv_bfloat16* dst = out + static_cast(row) * 2 * H + (is_hidden ? H : 0); #pragma unroll for (int it = 0; it < kIters; ++it) @@ -120,8 +136,11 @@ __global__ void __launch_bounds__(kThreads) int const c = (it * kThreads + threadIdx.x) * kVec; if (c < H) { - uint4 const wr = *reinterpret_cast(w + c); - __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + if constexpr (!kDraftInputPrewaitWeight) + { + wv[it] = *reinterpret_cast(w + c); + } + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wv[it]); uint4 o; __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); #pragma unroll diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu new file mode 100644 index 000000000000..a66c20a0c867 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu @@ -0,0 +1,267 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// The TP2 dense layers' QK norm + RoPE with FP8 output. The norm and RoPE are the generic fusedQKNormRopeKernel's at +// head_dim 128 without YaRN or mRoPE, expression for expression; the store rounds Q/K to BF16 and then to E4M3 +// (saturating), the two roundings of the stock BF16 kernel followed by its FP8 cast. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace denseQkNormRopeFp8 +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kVecSize = kElemsPerThread * sizeof(__nv_bfloat16) / 4; +constexpr int kBlockSize = 256; +static_assert(kVecSize == 2, "a lane loads 4 BF16 values as one uint2"); + +// tensorrt_llm::common::warpReduceSum: the same butterfly, so the same sum. +__device__ __forceinline__ float warpReduceSum(float val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + val = val + __shfl_xor_sync(0xffffffff, val, mask, 32); + } + return val; +} + +__device__ __forceinline__ void storeFp8ViaBf16(__nv_fp8_e4m3* out, long long offsetThread, float const (&elements)[4]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + float2 const rounded = __bfloat1622float2(__float22bfloat162_rn(make_float2(elements[i], elements[i + 1]))); + __nv_fp8x2_e4m3 const packed(rounded); + reinterpret_cast<__nv_fp8x2_storage_t*>(&out[offsetThread])[i / 2] = packed.__x; + } +} + +// The K (kv 0) or V (kv 1) row of one head of token tokenIdx in the HND paged pool, at the token's cache slot. +__device__ __forceinline__ __nv_fp8_e4m3* cacheRow(KvCacheView const& cache, int kv, int headIdx, int tokenIdx) +{ + long long const slot = static_cast(cache.outCacheLoc[tokenIdx]); + long long const page = slot / cache.tokensPerBlock; + long long const within = slot % cache.tokensPerBlock; + return cache.base + page * cache.pageStride + kv * cache.kvStride + headIdx * cache.headStride + + within * cache.tokenStride; +} + +template +__global__ void denseQkNormRopeFp8Kernel(__nv_bfloat16 const* qkvIn, __nv_fp8_e4m3* qkvOut, int numHeadsQ, + int numHeadsK, int numHeadsV, int rotaryDim, float eps, __nv_bfloat16 const* qWeight, __nv_bfloat16 const* kWeight, + float base, int const* positionIds, int numTokens, bool useGemma, KvCacheView cache) +{ + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId; + int const totalQkHeads = numHeadsQ + numHeadsK; + int const totalHeads = totalQkHeads + numHeadsV; + int const tokenIdx = globalWarpIdx / totalHeads; + int const localHeadIdx = globalWarpIdx % totalHeads; + if (tokenIdx >= numTokens) + { + return; + } + bool const isQ = localHeadIdx < numHeadsQ; + bool const isV = localHeadIdx >= totalQkHeads; + int headIdx; + int segStart; + if (isQ) + { + headIdx = localHeadIdx; + segStart = 0; + } + else if (!isV) + { + headIdx = localHeadIdx - numHeadsQ; + segStart = numHeadsQ * kHeadDim; + } + else + { + headIdx = localHeadIdx - totalQkHeads; + segStart = totalQkHeads * kHeadDim; + } + + using vec_T = uint2; + int const offsetThread + = tokenIdx * totalHeads * kHeadDim + segStart + headIdx * kHeadDim + laneId * kElemsPerThread; + float elements[kElemsPerThread]; + float sumOfSquares = 0.0f; + { + vec_T vec = *reinterpret_cast(&qkvIn[offsetThread]); +#pragma unroll + for (int i = 0; i < kVecSize; i++) + { + float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&vec) + i)); + sumOfSquares += vals.x * vals.x; + sumOfSquares += vals.y * vals.y; + elements[2 * i] = vals.x; + elements[2 * i + 1] = vals.y; + } + } + if (isV) + { + // V is BF16 already: its single E4M3 rounding is the stock cast's. + storeFp8ViaBf16(qkvOut, offsetThread, elements); + if (kWriteCache && tokenIdx < cache.numLiveTokens && cache.outCacheLoc[tokenIdx] >= 0) + { + storeFp8ViaBf16(cacheRow(cache, 1, headIdx, tokenIdx), laneId * kElemsPerThread, elements); + } + return; + } + + sumOfSquares = warpReduceSum(sumOfSquares); + float rms_rcp = rsqrtf(sumOfSquares / static_cast(kHeadDim) + eps); + for (int i = 0; i < kElemsPerThread; i++) + { + int dim = laneId * kElemsPerThread + i; + float weight = isQ ? __bfloat162float(qWeight[dim]) : __bfloat162float(kWeight[dim]); + elements[i] *= rms_rcp * (useGemma ? (1.0f + weight) : weight); + } + + float elements2[kElemsPerThread]; + float cos_vals[kElemsPerThread]; + float sin_vals[kElemsPerThread]; + float const neg2_log2base_over_rd = -2.0f * __log2f(base) / static_cast(rotaryDim); + int const rd_mask = rotaryDim - 1; + bool const rd_is_pow2 = ((rotaryDim & rd_mask) == 0); + float const pos_id = static_cast(positionIds[tokenIdx]); + if constexpr (kInterleave) + { + for (int i = 0; i < kElemsPerThread; i++) + { + if (i % 2 == 0) + { + elements2[i] = -elements[i + 1]; + } + else + { + elements2[i] = elements[i - 1]; + } + int dim_idx = laneId * kElemsPerThread + i; + int half_dim = dim_idx / 2; + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); + } + } + else + { + __syncwarp(); + int pairOffset = (rotaryDim / 2) / kElemsPerThread; + for (int i = 0; i < kElemsPerThread; i++) + { + elements2[i] = __shfl_xor_sync(0xffffffff, elements[i], pairOffset); + if (laneId < pairOffset) + { + elements2[i] = -elements2[i]; + } + int dim_idx = laneId * kElemsPerThread + i; + dim_idx = rd_is_pow2 ? ((dim_idx * 2) & rd_mask) : ((dim_idx * 2) % rotaryDim); + int half_dim = dim_idx / 2; + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); + } + __syncwarp(); + } + + // attention_factor is 1 without YaRN: (x * cos + x2 * sin) * 1 rounds as the generic kernel's product does. + float const attention_factor = 1.0f; + if (rotaryDim == kHeadDim) + { + for (int i = 0; i < kElemsPerThread; i++) + { + elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor; + } + } + else + { + for (int i = 0; i < kElemsPerThread; i++) + { + int dim_idx = laneId * kElemsPerThread + i; + if (dim_idx < rotaryDim) + { + elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor; + } + } + } + storeFp8ViaBf16(qkvOut, offsetThread, elements); + if constexpr (kWriteCache) + { + // Rows past the live count (padded graph rows) own no slot, nor does a negative (sentinel) slot. + if (!isQ && tokenIdx < cache.numLiveTokens && cache.outCacheLoc[tokenIdx] >= 0) + { + storeFp8ViaBf16(cacheRow(cache, 0, headIdx, tokenIdx), laneId * kElemsPerThread, elements); + } + // The compact query [tokens, q heads, 128] the decode reads, beside the packed row. + if (isQ && cache.qCompact != nullptr) + { + storeFp8ViaBf16(cache.qCompact, + (static_cast(tokenIdx) * numHeadsQ + headIdx) * kHeadDim + laneId * kElemsPerThread, + elements); + } + } +} + +} // namespace + +void launch(void const* qkvIn, void* qkvOut, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, int rotaryDim, + float eps, void const* qWeight, void const* kWeight, float base, bool interleave, int const* positionIds, + bool useGemma, KvCacheView const* cache, cudaStream_t stream) +{ + int const totalWarps = numTokens * (numHeadsQ + numHeadsK + numHeadsV); + int const warpsPerBlock = kBlockSize / 32; + dim3 const grid((totalWarps + warpsPerBlock - 1) / warpsPerBlock); + auto const* in = static_cast<__nv_bfloat16 const*>(qkvIn); + auto* out = static_cast<__nv_fp8_e4m3*>(qkvOut); + auto const* qw = static_cast<__nv_bfloat16 const*>(qWeight); + auto const* kw = static_cast<__nv_bfloat16 const*>(kWeight); + KvCacheView const view = cache != nullptr ? *cache : KvCacheView{}; + auto const run = [&](auto kernel) + { + kernel<<>>(in, out, numHeadsQ, numHeadsK, numHeadsV, rotaryDim, eps, qw, kw, base, + positionIds, numTokens, useGemma, view); + }; + if (cache != nullptr) + { + interleave ? run(denseQkNormRopeFp8Kernel) : run(denseQkNormRopeFp8Kernel); + } + else + { + interleave ? run(denseQkNormRopeFp8Kernel) : run(denseQkNormRopeFp8Kernel); + } +} + +} // namespace denseQkNormRopeFp8 +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h new file mode 100644 index 000000000000..ec94c1dbd550 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace denseQkNormRopeFp8 +{ + +constexpr int kHeadDim = 128; + +//! A layer's FP8 K/V pool in the HND paged layout [pages, 2 (K, V), kv heads, tokens per block, 128] (element +//! strides), with the step's per-token cache slots (page * tokensPerBlock + within) for its first numLiveTokens rows. +struct KvCacheView +{ + __nv_fp8_e4m3* base = nullptr; + int const* outCacheLoc = nullptr; + long long pageStride = 0; + long long kvStride = 0; + long long headStride = 0; + long long tokenStride = 0; + int tokensPerBlock = 1; + int numLiveTokens = 0; + //! Optional: every row's Q also stored compact, [tokens, q heads, 128]. + __nv_fp8_e4m3* qCompact = nullptr; +}; + +//! The MiniMax-M3 TP2 dense layers' per-head QK RMSNorm + partial RoPE of a packed BF16 [tokens, (q + k + v) * 128] +//! row, written out of place as FP8 E4M3 with V copy-cast. Q and K are rounded to BF16 before E4M3, as the stock path +//! (the BF16 fused_qk_norm_rope output cast to FP8) rounds them. With cache, the K and V rows of the live tokens are +//! also stored into the paged pool at their slots (the bytes of the qkvOut rows). +void launch(void const* qkvIn, void* qkvOut, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, int rotaryDim, + float eps, void const* qWeight, void const* kWeight, float base, bool interleave, int const* positionIds, + bool useGemma, KvCacheView const* cache, cudaStream_t stream); + +} // namespace denseQkNormRopeFp8 +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.cu new file mode 100644 index 000000000000..67e054e74712 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.cu @@ -0,0 +1,620 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h" + +#include +#include +#include + +// One pass of the fused greedy Eagle3 draft loop of MiniMax-M3 (GB300, TP4, <= 8 requests per pass). +// +// One kernel replaces the draft lm_head GEMM, the logits all-gather / float cast / argmax of the in-tree greedy draft +// sampling and the draft loop's per-pass bookkeeping, at the in-tree precision: +// +// * The vocab-sharded BF16 lm_head weight is stored in mma.m16n8k16 A-fragment order (the same values): per 16-row +// tile and k16 step, 512 bytes; lane (g, j) holds 16 bytes = {row g k 2j..2j+1, row g+8 k 2j.., row g k 2j+8.., +// row g+8 k 2j+8..}. +// * Or, by default, in a lossless 12-bit coding of those fragments (75% of the bytes; the target head's coding, see +// targetHead.cu): per 16-row tile and unit of 4 k16 steps a lane reads three 16-byte words, the sign/mantissa bytes +// of steps 0-1, of steps 2-3, then 4 code words (exponent - base, 4 bits; 15 = an exception). The decoded values +// are the checkpoint's BF16 bits; an exception decodes to +0 and the segment's flush adds the exact product x * w +// from the tile's list (products of BF16 values are exact in FP32). +// * A persistent grid (one 16-warp CTA per SM) streams whole tiles over the full K with BF16 MMAs and FP32 +// accumulation; the BF16 head input rows sit in shared memory in B-fragment order, unchanged. +// * Each logit is rounded to BF16 (the in-tree logits are the BF16 GEMM output) and becomes a 64-bit key +// ordered(value) << 32 | ~global_index, with -0 taken as +0: the larger value wins and ties pick the lowest index, as +// torch.max / torch.argmax over the gathered logits do. Keys are max-reduced per row: warp shuffles, then shared +// memory, then atomicMax. +// * The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping +// (Lamport: 0 = empty, a key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the +// slots, and runs the pass's bookkeeping. Call n uses slot half n % 2. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace draftStep +{ +namespace +{ +constexpr int kUnroll = 8; // k16 chunks in flight per warp +constexpr int kChunkBytes = 512; // one tile's k16 step: 32 lanes x 16 bytes +constexpr long long kMaxSpins = 1LL << 31; // cross-rank key poll bound + +// Coded units (4 k16 steps, 3 x 16 bytes per lane) in flight per warp: 3 up to 4 rows, 2 at 8 rows. +template +constexpr int codedUnroll() +{ + return MP >= 8 ? 2 : 3; +} + +__device__ __forceinline__ void mma_bf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// BF16 pair p (values 2p, 2p + 1) of a lane's 8 coded values: sign/mantissa bytes sm (4 per word), codes word cw. +__device__ __forceinline__ uint32_t coded_decode_pair(uint32_t sm_word, int byte0, uint32_t cw, int p, uint32_t base) +{ + uint32_t const b = __byte_perm(sm_word, 0u, byte0 == 0 ? 0x4140 : 0x4342); // {b0, 0, b1, 0} + uint32_t const q = (cw >> (4 * p)) & 0x000f000fu; + uint32_t const e = q * 128u + base; // exponent field of both halves + // Code 15 (an exception, sign/mantissa byte 0) decodes to +0. + return ((b & 0x007f007fu) | ((b << 8) & 0x80008000u) | e) & __vcmpne2(q, 0x000f000fu); +} + +// One coded unit (4 k16 steps) against the BF16 B fragments: steps alternate between the two accumulators, as the +// k16 steps of the plain BF16 head do. +template +__device__ __forceinline__ void chunk_mma_coded(uint4 const& sm01, uint4 const& sm23, uint4 const& cw, uint32_t base, + uint2 const* xs, int unit, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint32_t const smw[8] = {sm01.x, sm01.y, sm01.z, sm01.w, sm23.x, sm23.y, sm23.z, sm23.w}; + uint32_t const cws[4] = {cw.x, cw.y, cw.z, cw.w}; +#pragma unroll + for (int ci = 0; ci < 4; ++ci) + { + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[((unit * 4 + ci) * 4 + j) * MP + g]; + uint32_t const w0 = coded_decode_pair(smw[2 * ci], 0, cws[ci], 0, base); + uint32_t const w1 = coded_decode_pair(smw[2 * ci], 2, cws[ci], 1, base); + uint32_t const w2 = coded_decode_pair(smw[2 * ci + 1], 0, cws[ci], 2, base); + uint32_t const w3 = coded_decode_pair(smw[2 * ci + 1], 2, cws[ci], 3, base); + if (ci & 1) + mma_bf16(db, w0, w1, w2, w3, b.x, b.y); + else + mma_bf16(da, w0, w1, w2, w3, b.x, b.y); + } +} + +// Exceptions of `tile` with k in [kb, ke): v += x * w (exact) for the owner lanes (row g or g + 8; tokens 2j, +// 2j + 1 < M), x from the staged B fragments. Warp-uniform (shuffles). +template +__device__ __forceinline__ void coded_exceptions( + CodedArgs const& cd, int tile, int kb, int ke, uint2 const* xs, int M, int lane, float (&v)[4]) +{ + int const g = lane >> 2, j = lane & 3; + int const e0 = __ldg(cd.exc_off + tile), e1 = __ldg(cd.exc_off + tile + 1); + for (int b0 = e0; b0 < e1; b0 += 32) + { + uint32_t pos = 0xffffffffu; + float2 wv = make_float2(0.f, 0.f); + if (b0 + lane < e1) + { + pos = __ldg(cd.exc_pos + b0 + lane); + wv = __ldg(cd.exc_w + b0 + lane); + } + int const cnt = min(32, e1 - b0); + for (int q = 0; q < cnt; ++q) + { + uint32_t const pq = __shfl_sync(0xffffffffu, pos, q); + float const w0 = __shfl_sync(0xffffffffu, wv.x, q); + int const r = static_cast(pq >> 16), k = static_cast(pq & 0xffffu); + if (k >= kb && k < ke && (r & 7) == g) + { + int const st = k >> 4, rr = k & 15; + bool const hi = r >= 8; +#pragma unroll + for (int q2 = 0; q2 < 2; ++q2) + { + int const n = 2 * j + q2; + if (n < M) + { + uint2 const bw = xs[(st * 4 + ((rr & 7) >> 1)) * MP + n]; + uint32_t const word = rr >= 8 ? bw.y : bw.x; + float const xv = __uint_as_float((rr & 1) ? (word & 0xffff0000u) : (word << 16)); + float const d = __fmul_rn(xv, w0); + if (hi) + v[2 + q2] += d; + else + v[q2] += d; + } + } + } + __syncwarp(); + } + } +} + +__device__ __forceinline__ uint4 ldg_stream(uint4 const* p) +{ + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.u32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +// The argmax value of a logit: its BF16 rounding (the in-tree logits' dtype), -0 as +0. +__device__ __forceinline__ float key_value(float v) +{ + float const r = __bfloat162float(__float2bfloat16(v)); + return r == 0.f ? 0.f : r; +} + +__device__ __forceinline__ unsigned long long make_key(float v, uint32_t index) +{ + uint32_t u = __float_as_uint(v); + u = (u & 0x80000000u) ? ~u : (u | 0x80000000u); + return (static_cast(u) << 32) | static_cast(0xFFFFFFFFu - index); +} + +__device__ __forceinline__ unsigned long long shfl_max(unsigned long long k, int o) +{ + unsigned long long const ok = __shfl_xor_sync(0xffffffffu, k, o); + return ok > k ? ok : k; +} + +// One k16 step of a tile: the lane's A fragment w against the activation B fragment of step `chunk` (rows g < MP). +// Even and odd steps accumulate separately (two independent MMA chains); the tile's sum is da + db. +template +__device__ __forceinline__ void chunk_mma( + uint4 const& w, uint2 const* xs, int chunk, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[(chunk * 4 + j) * MP + g]; + if (chunk & 1) + mma_bf16(db, w.x, w.y, w.z, w.w, b.x, b.y); + else + mma_bf16(da, w.x, w.y, w.z, w.w, b.x, b.y); +} + +// Add a warp's segment of a tile into the CTA's accumulator (with a coded head, the segment's exceptions first). +template +__device__ __forceinline__ void flush_tile( + float4* acc, float (&da)[4], float (&db)[4], Params const& p, int tile, int kb, int ke, uint2 const* xs, int lane) +{ + float v[4] = {da[0] + db[0], da[1] + db[1], da[2] + db[2], da[3] + db[3]}; + if constexpr (CODED) + coded_exceptions(p.cd, tile, kb, ke, xs, p.M, lane, v); + float* a = reinterpret_cast(acc); + atomicAdd(a + 0, v[0]); + atomicAdd(a + 1, v[1]); + atomicAdd(a + 2, v[2]); + atomicAdd(a + 3, v[3]); +#pragma unroll + for (int e = 0; e < 4; ++e) + { + da[e] = 0.f; + db[e] = 0.f; + } +} + +// Balanced streaming GEMV over the fragment-ordered weight, persistent grid of 16-warp CTAs. CTA c owns tiles +// [t_begin, t_end); its (tile, chunk) space is split into 16 equal contiguous ranges, one per warp, so every warp +// streams the same number of bytes. A tile straddles at most two warps; their partial sums are added into shared +// memory (two addends: the order does not change the sum). MP = rows padded to 1 / 2 / 4 / 8. A chunk is one k16 +// step of the BF16 head, or one coded unit (4 k16 steps) of the coded head (CODED: no shared-memory prologue). +template +__global__ void __launch_bounds__(kWarps * 32, 1) + draft_step_kernel(Params const p, int kchunks, int pf_chunks, int pf_off) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(16) uint2 xs[]; // [K/16 steps][4 (j)][MP (g)], then the weight prologue + __shared__ float4 acc_s[kMaxTilesPerCta + 1][32]; // + 1: an equal-chunk span can touch one more tile + __shared__ unsigned long long red[kWarps][kMaxTokens]; + __shared__ __align__(8) unsigned long long pf_bar[kWarps]; + __shared__ bool s_last; + __shared__ int s_tok[kMaxTokens]; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const g = lane >> 2; + int const j = lane & 3; + constexpr int kChunkK = CODED ? 64 : 16; // k per chunk + constexpr int kWords = CODED ? 3 : 1; // 16-byte words per lane and chunk + constexpr int KU = CODED ? codedUnroll() : kUnroll; + int const K = kchunks * kChunkK; + int const M = p.M; + int const r_lo = p.tile_lo; + int const r_n = (p.tile_hi > 0 ? p.tile_hi : p.n_tiles) - r_lo; + bool const eq = p.bnd != nullptr; + // Whole tiles: CTA c owns tiles [t_begin, t_end). Equal chunks: CTA c streams chunks [c0, c1) of the range, which + // start off0 chunks into tile t_begin (its first and last tiles may be shared with the neighbouring CTAs). + long long const r_chunks = static_cast(r_n) * kchunks; + long long const c0 = eq ? (static_cast(blockIdx.x) * r_chunks) / gridDim.x : 0; + long long const c1 = eq ? (static_cast(blockIdx.x + 1) * r_chunks) / gridDim.x : 0; + int const t_begin = r_lo + + (eq ? static_cast(c0 / kchunks) + : static_cast((static_cast(blockIdx.x) * r_n) / gridDim.x)); + int const t_end = r_lo + + (eq ? static_cast((c1 + kchunks - 1) / kchunks) + : static_cast((static_cast(blockIdx.x + 1) * r_n) / gridDim.x)); + int const n_local = t_end - t_begin; + int const off0 = eq ? static_cast(c0 - static_cast(t_begin - r_lo) * kchunks) : 0; + int const span = eq ? static_cast(c1 - c0) : n_local * kchunks; + int const q_begin = off0 + warp * span / kWarps; + int const q_end = off0 + (warp + 1) * span / kWarps; + // Equal chunks: whether the first / last local tile is shared with the previous / next CTA of the launch. + bool const shared_first = eq && off0 != 0; + bool const shared_last = eq && ((c1 % kchunks) != 0); + long long const base = static_cast(t_begin) * kchunks; // first chunk of this CTA + uint4 const* const wf = p.wf; + + // Weights are static: start this warp's stream before waiting on the producer of x. The warp's first pf_chunks + // chunks come in by one bulk copy into shared memory (dynamic offset pf_off), which fills the grid-dependency + // wait and the staging; the register ring starts after them. + int const n_pf = CODED ? 0 : min(pf_chunks, q_end - q_begin); + uint4* const wpf = reinterpret_cast(reinterpret_cast(xs) + pf_off) + warp * pf_chunks * 32; + if (n_pf > 0 && lane == 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + uint32_t const bytes = static_cast(n_pf) * kChunkBytes; + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + static_cast(__cvta_generic_to_shared(wpf))), + "l"(wf + (base + q_begin) * 32), "r"(bytes), "r"(bar) + : "memory"); + } + __syncwarp(); + int const q_reg = q_begin + n_pf; + uint4 buf[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q_reg + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = ldg_stream(wf + ((base + q_reg + u) * kWords + l) * 32 + lane); + } + } + for (int i = threadIdx.x; i < n_local * 32; i += blockDim.x) + acc_s[i / 32][i % 32] = make_float4(0.f, 0.f, 0.f, 0.f); + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // Stage the head input rows in B-fragment order: item (k16 step S, row gg) = two 16-byte loads of x, four + // fragment words {x[2jj], x[2jj + 1]}, {x[2jj + 8], x[2jj + 9]} (BF16 pairs, moved unchanged). Rows >= M are 0. + constexpr int kSU = MP >= 4 ? 4 : 2; + int const n_items = (K / 16) * MP; + for (int i0 = threadIdx.x; i0 < n_items; i0 += kSU * blockDim.x) + { + uint4 xv[kSU][2]; +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + int const gg = i % MP; + int const S = i / MP; + xv[u][0] = make_uint4(0u, 0u, 0u, 0u); + xv[u][1] = make_uint4(0u, 0u, 0u, 0u); + if (i < n_items && gg < M) + { + long long const xr = p.x_rows != nullptr ? p.x_rows[gg] : gg; + uint4 const* xp = reinterpret_cast(p.x + xr * K + 16 * S); + xv[u][0] = xp[0]; + xv[u][1] = xp[1]; + } + } +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + if (i >= n_items) + continue; + int const gg = i % MP; + int const S = i / MP; + uint32_t const lo[4] = {xv[u][0].x, xv[u][0].y, xv[u][0].z, xv[u][0].w}; + uint32_t const hi[4] = {xv[u][1].x, xv[u][1].y, xv[u][1].z, xv[u][1].w}; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) + xs[(S * 4 + jj) * MP + gg] = make_uint2(lo[jj], hi[jj]); + } + } + __syncthreads(); + + bool const b_active = g < MP; + float da[4] = {0.f, 0.f, 0.f, 0.f}; + float db[4] = {0.f, 0.f, 0.f, 0.f}; + int tl = q_begin / kchunks; // local tile of the current accumulator + int c = q_begin - tl * kchunks; // chunk within that tile + int cs = c; // first chunk of the warp's segment of that tile + if (n_pf > 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + asm volatile( + "{\n .reg .pred p;\n PF_WAIT%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n" + " @!p bra PF_WAIT%=;\n}\n" ::"r"(bar) + : "memory"); + for (int q = q_begin; q < q_reg; ++q) + { + chunk_mma(wpf[(q - q_begin) * 32 + lane], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane); + c = 0; + cs = 0; + ++tl; + } + } + } + for (int q = q_reg; q < q_end; q += KU) + { + uint4 nxt[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + KU + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + nxt[u][l] = ldg_stream(wf + ((base + q + KU + u) * kWords + l) * 32 + lane); + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + u < q_end) + { + if constexpr (CODED) + chunk_mma_coded(buf[u][0], buf[u][kWords > 1 ? 1 : 0], buf[u][kWords > 2 ? 2 : 0], p.cd.base, + xs, c, j, g, b_active, da, db); + else + chunk_mma(buf[u][0], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane); + c = 0; + cs = 0; + ++tl; + } + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = nxt[u][l]; + } + if (c != 0) + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, c * kChunkK, xs, lane); + __syncthreads(); + + // Accumulator lane (g, j) of a tile: rows g / g + 8 (x, y / z, w) of tokens 2j / 2j + 1 (x, z / y, w). + unsigned long long best0 = 0ull, best1 = 0ull; // tokens 2j, 2j + 1 + for (int t = warp; t < n_local; t += kWarps) + { + float4 v = acc_s[t][lane]; + if ((t == 0 && shared_first) || (t == n_local - 1 && shared_last)) + { + // A tile shared with a neighbouring CTA: add this CTA's partial into bnd; the second CTA to arrive reads + // the sum (two addends), keys the tile and restores bnd and its counter to zero for the next call. + int const tile = t_begin + t; + float* const bp = reinterpret_cast(p.bnd + static_cast(tile) * 32 + lane); + atomicAdd(bp + 0, v.x); + atomicAdd(bp + 1, v.y); + atomicAdd(bp + 2, v.z); + atomicAdd(bp + 3, v.w); + __threadfence(); + __syncwarp(); + unsigned int prev = 0; + if (lane == 0) + prev = atomicAdd(p.bnd_cnt + tile, 1u); + prev = __shfl_sync(0xffffffffu, prev, 0); + if (prev == 0) + continue; // the neighbour keys it + __threadfence(); + v = __ldcg(p.bnd + static_cast(tile) * 32 + lane); + p.bnd[static_cast(tile) * 32 + lane] = make_float4(0.f, 0.f, 0.f, 0.f); + if (lane == 0) + p.bnd_cnt[tile] = 0u; + } + int const r0 = (t_begin + t) * 16 + g; + uint32_t const i0 = static_cast(p.vocab_offset + r0); + unsigned long long k0 = make_key(key_value(v.x), i0); + unsigned long long const k0b = make_key(key_value(v.z), i0 + 8); + unsigned long long k1 = make_key(key_value(v.y), i0); + unsigned long long const k1b = make_key(key_value(v.w), i0 + 8); + k0 = k0b > k0 ? k0b : k0; + k1 = k1b > k1 ? k1b : k1; + best0 = k0 > best0 ? k0 : best0; + best1 = k1 > best1 ? k1 : best1; + } +#pragma unroll + for (int o = 4; o < 32; o <<= 1) + { + best0 = shfl_max(best0, o); + best1 = shfl_max(best1, o); + } + asm volatile("griddepcontrol.launch_dependents;"); + if (lane < 4) + { + red[warp][2 * lane] = best0; + red[warp][2 * lane + 1] = best1; + } + __syncthreads(); + if (threadIdx.x < M) + { + unsigned long long b = 0ull; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + b = red[w][threadIdx.x] > b ? red[w][threadIdx.x] : b; + if (b != 0ull) + atomicMax(p.keys + threadIdx.x, b); + } + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned int const prev = atomicAdd(p.counters, 1u); + s_last = prev == (p.total_ctas != 0u ? p.total_ctas : gridDim.x) - 1u; + } + __syncthreads(); + if (!s_last) + return; + + __threadfence(); + unsigned int const call = *reinterpret_cast(p.counters + 1); + int const half = static_cast(call & 1u); + if (threadIdx.x < M) + { + int const t = threadIdx.x; + unsigned long long const mine = atomicExch(p.keys + t, 0ull); + unsigned long long* slot_mc = p.ws_mc + (half * kMaxTokens + t) * p.world; + unsigned long long* slot_uc = p.ws_uc + (half * kMaxTokens + t) * p.world; + slot_mc[p.rank] = mine; + unsigned long long best = 0ull; + for (int r = 0; r < p.world; ++r) + { + unsigned long long v; + long long spins = 0; + do + { + v = *reinterpret_cast(slot_uc + r); + if (++spins > kMaxSpins) + __trap(); // a missing peer key becomes a CUDA error, never a hang (the mc all-reduce's bound) + } while (v == 0ull); + best = v > best ? v : best; + slot_uc[r] = 0ull; + } + s_tok[t] = static_cast(0xFFFFFFFFu - static_cast(best & 0xFFFFFFFFull)); + if (p.out_keys != nullptr) + p.out_keys[t] = best; + } + __syncthreads(); + if (threadIdx.x == 0) + { + p.counters[0] = 0u; + p.counters[1] = call + 1u; + } + if (threadIdx.x < M) + { + StepArgs const& st = p.st; + int const b = threadIdx.x; + int const tok = s_tok[b]; + st.draft_out[b * st.draft_stride + st.step] = tok; + st.ids_next[b] = tok; + long long const r = p.x_rows != nullptr ? p.x_rows[b] : b; + st.pos_next[b] = st.pos_src[r] + 1; + if (st.step == 0 && st.num_acc != nullptr) + { + st.kv_lens[b] -= st.draft_len - st.num_acc[b]; + st.seq_lens[b] = 1; + } + else if (st.step > 0) + { + st.kv_lens[b] += 1; + } + } +#endif +} + +template +cudaError_t launch(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + auto kern = draft_step_kernel; + size_t const xs_bytes = static_cast(K / 16) * 4 * MP * sizeof(uint2); + // Shared-memory weight prologue: chunks per warp that fit next to the kernel's static + dynamic shared memory. + static int const sStaticSmem = [] + { + cudaFuncAttributes fa{}; + return cudaFuncGetAttributes(&fa, draft_step_kernel) == cudaSuccess + ? static_cast(fa.sharedSizeBytes) + : -1; + }(); + static int const sMaxSmemPerBlock = [] + { + int dev = 0, maxSmem = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + return maxSmem; + }(); + int pf = 0; + if (!CODED && sStaticSmem >= 0) + { + long long const room + = static_cast(sMaxSmemPerBlock) - sStaticSmem - static_cast(xs_bytes) - 1024; + long long const want = static_cast(kPrefetchKb) * 1024; + pf = static_cast(std::max(0ll, std::min(room, want)) / (kWarps * kChunkBytes)); + } + size_t const smem = xs_bytes + static_cast(pf) * kWarps * kChunkBytes; + cudaError_t err = cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + return err; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(grid)); + cfg.blockDim = dim3(kWarps * 32); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, kern, p, K / (CODED ? 64 : 16), pf, static_cast(xs_bytes)); +} + +template +cudaError_t dispatchRows(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + if (p.M <= 1) + return launch<1, CODED>(p, K, grid, pdl, stream); + if (p.M <= 2) + return launch<2, CODED>(p, K, grid, pdl, stream); + if (p.M <= 4) + return launch<4, CODED>(p, K, grid, pdl, stream); + return launch<8, CODED>(p, K, grid, pdl, stream); +} + +} // namespace + +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + bool const coded = p.cd.exc_off != nullptr; + if (p.M < 1 || p.M > kMaxTokens || K % (coded ? 64 : 16) != 0 || grid < 1) + return cudaErrorInvalidValue; + return coded ? dispatchRows(p, K, grid, pdl, stream) : dispatchRows(p, K, grid, pdl, stream); +} + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h new file mode 100644 index 000000000000..eb38b73488a4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +//! One pass of the fused greedy Eagle3 draft loop of the MiniMax-M3 NVFP4 TP4 decode path (sm_103a): the BF16 +//! vocab-sharded draft head, the cross-rank argmax and the draft loop's per-pass bookkeeping in one kernel. See +//! draftStep.cu. +namespace draftStep +{ + +constexpr int kMaxTokens = 8; // head rows (requests) of one launch +constexpr int kWarps = 16; // warps per CTA +constexpr int kMaxTilesPerCta = 32; // 16-row weight tiles one CTA can own +//! Shared-memory weight prologue per CTA, in KB (capped by the free shared memory). +constexpr int kPrefetchKb = 112; + +// Draft-pass epilogue (last CTA, one thread per request row b): stores the token as draft token `step`, stages the +// next pass's input id and position, and advances the draft KV lengths the way the in-tree draft loop does between +// passes (pass 0 with num_acc: drop the rejected verify rows and switch to one token per request; later passes: +1). +struct StepArgs +{ + int* draft_out; // [nreq, draft_stride]: column `step` receives the token + int draft_stride; + int step; + int* ids_next; // [nreq] + int const* pos_src; // positions of the rows the head read (indexed like x) + int* pos_next; // [nreq] = pos_src[x_rows[b] or b] + 1 + int* kv_lens; // [nreq] draft KV lengths, updated in place + int const* num_acc; // pass 0: accepted counts (nullptr: no rollback) + int* seq_lens; // pass 0: set to 1 + int draft_len; +}; + +// The 12-bit coded head (see draftStep.cu): the per-tile exception lists and the exponent base. exc_off == nullptr: +// the head is the plain BF16 fragment-ordered copy. +struct CodedArgs +{ + int const* exc_off; // [n_tiles + 1] exception list offsets + uint32_t const* exc_pos; // row in tile << 16 | k, per exception + float2 const* exc_w; // {w, 0} per exception: the exact weight + uint32_t base; // exponent base << 7 in both BF16 halves (base << 7 | base << 23) +}; + +// Everything a draft-step launch takes (see invokeDraftStep). +struct Params +{ + __nv_bfloat16 const* x; // [*, K] head input rows + long long const* x_rows; // optional [M]: head input row of request b (pass 0: the last accepted verify row) + int M; // requests, <= kMaxTokens + uint4 const* wf; // [n_tiles][K / 16][32 lanes] 16-byte BF16 A fragments, or the 12-bit coded stream + // [n_tiles][K / 64][3][32 lanes] when cd.exc_off is set (see draftStep.cu) + int n_tiles; + int vocab_offset; // global id of this rank's first head row + unsigned long long* keys; + unsigned int* counters; + unsigned long long* ws_uc; // this rank's exchange slots [2][kMaxTokens][world] + unsigned long long* ws_mc; // their multicast address + int rank; + int world; + unsigned long long* out_keys; // optional [M]: the winning global key of each row + StepArgs st; + CodedArgs cd; + // A pass split over several launches (one per locality domain, each on that domain's SMs): this launch streams + // tiles [tile_lo, tile_hi) (tile_hi 0 = n_tiles), and the last of total_ctas CTAs over all of the pass's launches + // runs the exchange and the bookkeeping (0 = this launch's grid: one launch per pass). + int tile_lo = 0; + int tile_hi = 0; + unsigned int total_ctas = 0; + // Equal-chunk split (set both, zero-filled [n_tiles][32] float4 / [n_tiles], restored to zero by the kernel): each + // CTA of a launch streams an equal share of its range's (tile, chunk) space instead of whole tiles, so a tile at a + // CTA boundary gets one partial sum from each of its two CTAs, added in bnd (two addends: the order does not change + // the sum) and keyed by the second to arrive (bnd_cnt). nullptr: whole tiles per CTA. + float4* bnd = nullptr; + unsigned int* bnd_cnt = nullptr; +}; + +//! Launch the draft step over `grid` 512-thread CTAs (K = the hidden size, a multiple of 16; of 64 for a coded head). +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream); + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu new file mode 100644 index 000000000000..179770d14a0e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu @@ -0,0 +1,189 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// The TP2 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's GPT-NeoX +// path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin negated, +// dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale E4M3 +// store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query +// saturates the same way (448, not NaN). + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace drafterRopeFp8Kv +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kBlockSize = 256; + +template +__device__ __forceinline__ float rotate(float c, float x, float s, float pair) +{ + if constexpr (kRopeMode == 0) + { + return __fadd_rn(__fmul_rn(c, x), __fmul_rn(s, pair)); + } + else if constexpr (kRopeMode == 1) + { + return __fmaf_rn(c, x, __fmul_rn(s, pair)); + } + else + { + return __fmaf_rn(s, pair, __fmul_rn(c, x)); + } +} + +__device__ __forceinline__ void storeE4m3Pairs(__nv_fp8_e4m3* dst, __nv_bfloat16 const (&v)[kElemsPerThread]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + __nv_bfloat162 pair; + pair.x = v[i]; + pair.y = v[i + 1]; + __nv_fp8x2_e4m3 const packed(pair); + reinterpret_cast<__nv_fp8x2_storage_t*>(dst)[i / 2] = packed.__x; + } +} + +template +__global__ void __launch_bounds__(kBlockSize) drafterRopeFp8KvKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const lane = threadIdx.x % 32; + int const warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; + int const slots = p.numHeadsQ + p.numHeadsKv; + int const token = warp / slots; + if (token >= p.numTokens) + { + return; + } + int const slot = warp % slots; + bool const isQ = slot < p.numHeadsQ; + int const head = isQ ? slot : slot - p.numHeadsQ; + int const request = token / p.qLen; + int const row = token % p.qLen; + int const cacheLen = p.seqLens[request]; + int const kvPos = cacheLen - p.qLen + row; + int const ropePos + = p.specPositionOffsets != nullptr ? p.specPositionOffsets[request * p.qLen + row] + cacheLen - p.qLen : kvPos; + int const rowWidth = (p.numHeadsQ + 2 * p.numHeadsKv) * kHeadDim; + __nv_bfloat16 const* src + = p.qkv + static_cast(token) * rowWidth + (isQ ? head * kHeadDim : (p.numHeadsQ + head) * kHeadDim); + int const half = p.rotaryDim / 2; + float2 const* coef = p.cosSin + static_cast(ropePos) * half; + + __nv_bfloat16 rotated[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + int const d = lane * kElemsPerThread + i; + bool const firstHalf = d < half; + float const x = __bfloat162float(src[d]); + float const pair = __bfloat162float(src[firstHalf ? d + half : d - half]); + float2 const cs = d < p.rotaryDim ? coef[firstHalf ? d : d - half] : make_float2(1.0f, 0.0f); + float const s = firstHalf ? -cs.y : cs.y; + rotated[i] = __float2bfloat16(rotate(cs.x, x, s, pair)); + } + + if (isQ) + { + storeE4m3Pairs( + p.qOut + (static_cast(token) * p.numHeadsQ + head) * kHeadDim + lane * kElemsPerThread, rotated); + return; + } + + int const block = kvPos / p.tokensPerBlock; + long long const inPage = (static_cast(head) * p.tokensPerBlock + kvPos % p.tokensPerBlock) * kHeadDim + + lane * kElemsPerThread; + long long const pageElems = static_cast(p.numHeadsKv) * p.tokensPerBlock * kHeadDim; + int const* table = p.blockTable + request * p.blockTableRequestStride; + long long const kPage = table[block]; + long long const vPage = table[p.blockTableKvStride + block]; + storeE4m3Pairs(p.pool + kPage * pageElems + inPage, rotated); + + __nv_bfloat16 const* vSrc = p.qkv + static_cast(token) * rowWidth + + (p.numHeadsQ + p.numHeadsKv + head) * kHeadDim + lane * kElemsPerThread; + __nv_bfloat16 v[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + v[i] = vSrc[i]; + } + storeE4m3Pairs(p.pool + vPage * pageElems + inPage, v); +} + +//! With PDL and maxSharedCarveout, the maximum shared-memory carveout lets the full-context decode's large-smem CTAs +//! (the dependent this kernel releases right after its wait) co-reside on the SMs this small kernel occupies. Set once +//! per instantiation. +template +cudaError_t launchWithMaxSharedCarveout( + cudaLaunchConfig_t const& config, Params const& params, bool pdl, bool maxSharedCarveout) +{ + auto const kernel = drafterRopeFp8KvKernel; + if (pdl && maxSharedCarveout) + { + static cudaError_t const carveout = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (carveout != cudaSuccess) + { + return carveout; + } + } + return cudaLaunchKernelEx(&config, kernel, params); +} + +} // namespace + +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream) +{ + int const warps = params.numTokens * (params.numHeadsQ + params.numHeadsKv); + int const warpsPerBlock = kBlockSize / 32; + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((warps + warpsPerBlock - 1) / warpsPerBlock); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + switch (params.ropeMode) + { + case 0: return launchWithMaxSharedCarveout<0>(config, params, pdl, maxSharedCarveout); + case 1: return launchWithMaxSharedCarveout<1>(config, params, pdl, maxSharedCarveout); + case 2: return launchWithMaxSharedCarveout<2>(config, params, pdl, maxSharedCarveout); + default: return cudaErrorInvalidValue; + } +} + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h new file mode 100644 index 000000000000..3302c94390f3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace drafterRopeFp8Kv +{ + +constexpr int kHeadDim = 128; + +struct Params +{ + //! [tokens, (numHeadsQ + 2 * numHeadsKv) * 128] BF16: the drafter layer's packed Q, K, V rows. + __nv_bfloat16 const* qkv; + //! [positions, rotaryDim / 2] (cos, sin) pairs: the layer's rotary_cos_sin table. + float2 const* cosSin; + //! [requests] cache lengths including this step's rows. + int const* seqLens; + //! [requests, qLen] spec-decoding position offsets of a multi-token pass, or null. + int const* specPositionOffsets; + //! [tokens, numHeadsQ, 128] E4M3 query out. + __nv_fp8_e4m3* qOut; + //! Base of the layer's K / V sub-pages, [*, numHeadsKv, tokensPerBlock, 128] E4M3. + __nv_fp8_e4m3* pool; + //! [requests, 2, maxBlocks] sub-page indices: K in row 0, V in row 1. + int const* blockTable; + long long blockTableRequestStride; + long long blockTableKvStride; + int numTokens; + int qLen; + int numHeadsQ; + int numHeadsKv; + int rotaryDim; + int tokensPerBlock; + //! How the rotation's two products are summed: 0 two products then the add, 1 fma(cos, x, sin * pair), + //! 2 fma(sin, pair, cos * x). + int ropeMode; +}; + +//! The MiniMax-M3 TP2 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, +//! rounded to BF16 (the in-tree preprocess's q output and cache input), then Q, K and V to saturating E4M3: K / V +//! as the in-tree preprocess stores them at unit KV scale, Q as torch's CUDA cast of the BF16 query. Rows map to +//! (request token / qLen, row token % qLen) at cache position seqLen - qLen + row. +//! +//! PDL (pdl): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction and the kernel +//! releases its dependent right after it. The dependent may read only step-invariant data before its own wait. +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream); + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu new file mode 100644 index 000000000000..dd0ad05437a6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu @@ -0,0 +1,86 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +// Row i of the step (lane c = column), as the Triton fold tail's program i: the store's padded rows, then the next +// step's inputs read from them (zero past each source width). +__global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) +{ + int const i = blockIdx.x; + int const c = threadIdx.x; + if (i >= *p.ctrl) + { + return; + } + int64_t const s = p.slots[i]; + int32_t const nw = c < p.snW ? p.srcNew[i * p.sn0 + c * p.sn1] : 0; + if (c < p.dnW) + { + p.dstNew[s * p.dn0 + c * p.dn1] = nw; + } + int32_t const nxt = c < p.sxW ? p.srcNext[i * p.sx0 + c * p.sx1] : 0; + if (c < p.dxW) + { + p.dstNext[s * p.dx0 + c * p.dx1] = nxt; + } + int32_t const drf = c < p.sdW ? p.srcDraft[i * p.sd0 + c * p.sd1] : 0; + if (c < p.ddW) + { + p.dstDraft[s * p.dd0 + c * p.dd1] = drf; + } + int32_t const ln = p.srcLens[i]; + int32_t const t = p.tokensPerStep; + int32_t const d = t - 1; + if (c == 0) + { + p.dstLens[s] = ln; + p.kvOff[i] = ln - t; + } + if (c < t) + { + p.ids[i * t + c] = c < p.dxW ? nxt : 0; + p.posOff[i * t + c] = ln; + } + if (c < d) + { + p.draftIds[i * d + c] = c < p.ddW ? drf : 0; + } +} + +} // namespace + +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream) +{ + foldTailKernel<<>>(p); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h new file mode 100644 index 000000000000..929ac2fed9ed --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h @@ -0,0 +1,72 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +//! Operands of the speculative step's forward-tail store and next-step staging (kernels/m3_fold_tail.py). Every tensor +//! is int32 except `slots` (int64); strides are in elements. The store views are [slots, width], the sources +//! [rows, width]. +struct FoldTailParams +{ + int32_t const* ctrl; // ctrl[0]: the rows the step armed (0: write nothing) + int64_t const* slots; + int32_t const* srcNew; + int64_t sn0, sn1; + int32_t snW; + int32_t* dstNew; + int64_t dn0, dn1; + int32_t dnW; + int32_t const* srcNext; + int64_t sx0, sx1; + int32_t sxW; + int32_t* dstNext; + int64_t dx0, dx1; + int32_t dxW; + int32_t const* srcDraft; + int64_t sd0, sd1; + int32_t sdW; + int32_t* dstDraft; + int64_t dd0, dd1; + int32_t ddW; + int32_t const* srcLens; + int32_t* dstLens; + int32_t* ids; // [rows * T] + int32_t* draftIds; // [rows * (T - 1)] + int32_t* posOff; // [rows * T] + int32_t* kvOff; // [rows] + int32_t tokensPerStep; +}; + +//! One CTA of 32 lanes per row, `rows` CTAs; the first ctrl[0] rows write. Widths and tokensPerStep must be <= 32. +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu new file mode 100644 index 000000000000..236a13e66d1f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu @@ -0,0 +1,223 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// The split-KV merge for the MiniMax-M3 NVFP4 TP2 (Rubin) full-context decode. The producer is the PrimTS split-KV +// decode with its separate-reduction scratch: per (request, KV head, split, row) one normalized 16-bit partial O and +// one fp32 log2-LSE. Active splits per row follow the producer's split ranges (the causal union of the row's grouped +// Q CTA, instruction-group aligned): splits past them hold nothing and are not read. +// +// One CTA per (request, KV head, kRowsPerCta rows); per row kGroups split groups x kThreadsPerRow threads x 8 dims. +// A thread issues the loads of every split it owns (strided over the groups) before using any, folds them under its +// group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to +// BF16 once. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace fullCtxReduce +{ +namespace +{ + +constexpr int kGroups = 8; +constexpr int kThreadsPerRow = kHeadDim / 8; +constexpr int kRowsPerCta = 2; +constexpr int kSplitsPerGroup = kMaxSplits / kGroups; +constexpr int kThreads = kGroups * kThreadsPerRow * kRowsPerCta; + +__device__ int activeSplits(Params const& p, int b, int row) +{ + int const seqKv = p.seqLensKv[b]; + int visible = seqKv; + if (p.causal) + { + int const token = row / p.headsPerKv; + int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; + int const end = min(base + p.qTokensPerCta, p.seqLenQ); + visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); + } + int const tiles = (visible + p.tileKv - 1) / p.tileKv; + int const tilesPerGroup = p.splits * p.instsKv; + int const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; + int const local = max(groups * p.instsKv, p.instsKv); + return (tiles + local - 1) / local; +} + +template +__device__ float2 unpack(uint32_t word) +{ + if constexpr (kBf16Partial) + { + return __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&word)); + } + else + { + return __half22float2(*reinterpret_cast<__half2 const*>(&word)); + } +} + +template +__global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const b = blockIdx.z; + int const h = blockIdx.y; + int const localRow = threadIdx.x / (kGroups * kThreadsPerRow); + int const row = blockIdx.x * kRowsPerCta + localRow; + int const g = (threadIdx.x / kThreadsPerRow) % kGroups; + int const t = threadIdx.x % kThreadsPerRow; + bool const validRow = row < p.rows; + int const active = validRow ? activeSplits(p, b, row) : 0; + + __shared__ float sMax[kRowsPerCta][kGroups]; + __shared__ float sSum[kRowsPerCta][kGroups]; + __shared__ float sAcc[kRowsPerCta][kGroups][kHeadDim]; + + uint4 const* partialO = static_cast(p.partialO); + size_t const kvIdx = static_cast(b) * p.numKvHeads + h; + float lse[kSplitsPerGroup]; + uint4 vec[kSplitsPerGroup]; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + int const s = g + i * kGroups; + if (s < active) + { + size_t const r = (kvIdx * p.splits + s) * p.rows + row; + lse[i] = __ldcg(p.partialStats + r); + vec[i] = __ldcg(partialO + r * kThreadsPerRow + t); + } + } + float groupMax = -CUDART_INF_F; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active) + { + groupMax = fmaxf(groupMax, lse[i]); + } + } + // A split of the CTA's causal union can hold no key of this row (log2-LSE -inf). Such a split, and a group of + // only such splits, is skipped rather than weighted 0: its partial O carries no defined value. + float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + float sum = 0.f; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active && lse[i] != -CUDART_INF_F) + { + float const w = exp2f(lse[i] - groupMax); + sum += w; + uint32_t const* words = reinterpret_cast(&vec[i]); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = unpack(words[j]); + acc[2 * j] += w * f.x; + acc[2 * j + 1] += w * f.y; + } + } + } + if (t == 0) + { + sMax[localRow][g] = groupMax; + sSum[localRow][g] = sum; + } +#pragma unroll + for (int j = 0; j < 8; ++j) + { + sAcc[localRow][g][t * 8 + j] = acc[j]; + } + __syncthreads(); + if (g != 0 || !validRow) + { + return; + } + float rowMax = -CUDART_INF_F; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + rowMax = fmaxf(rowMax, sMax[localRow][k]); + } + float total = 0.f; + float o[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + if (sMax[localRow][k] == -CUDART_INF_F) + { + continue; + } + float const scale = exp2f(sMax[localRow][k] - rowMax); + total += scale * sSum[localRow][k]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + o[j] += scale * sAcc[localRow][k][t * 8 + j]; + } + } + float const inv = total > 0.f ? 1.f / total : 0.f; + int const token = row / p.headsPerKv; + int const head = row % p.headsPerKv; + __nv_bfloat162 packed[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + packed[j] = __floats2bfloat162_rn(o[2 * j] * inv, o[2 * j + 1] * inv); + } + size_t const outRow + = (static_cast(b) * p.seqLenQ + token) * (p.numKvHeads * p.headsPerKv) + h * p.headsPerKv + head; + reinterpret_cast(p.out + outRow * kHeadDim)[t] = *reinterpret_cast(packed); +} + +} // namespace + +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) +{ + if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((params.rows + kRowsPerCta - 1) / kRowsPerCta, params.numKvHeads, params.batch); + config.blockDim = dim3(kThreads); + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + return params.bf16Partial ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) + : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); +} + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h new file mode 100644 index 000000000000..eca13e5d7dab --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +//! The split-KV merge of the MiniMax-M3 TP2 full-context decode (PrimTS split-KV decode, separate-reduction scratch) +//! straight to BF16. See fullCtxReduce.cu. +namespace fullCtxReduce +{ + +constexpr int kHeadDim = 128; +constexpr int kMaxSplits = 128; + +struct Params +{ + void const* partialO; // [B, Hkv, S, R, 128] normalized 16-bit partial O per split + float const* partialStats; // [B, Hkv, S, R] log2-LSE per split + int const* seqLensKv; // [B] + __nv_bfloat16* out; // [B, SQ, Hkv * hR, 128] + int batch; + int numKvHeads; + int splits; // S, the planned split capacity (<= kMaxSplits) + int rows; // R = hR * SQ, rows [token, local head] + int seqLenQ; // SQ + int headsPerKv; // hR + int qTokensPerCta; // query tokens one producer CTA owns (grouped Q) + int tileKv; // producer K/V tile + int instsKv; // producer K/V instructions per loop step + bool causal; + bool bf16Partial; +}; + +//! Merges the active splits of every row in fp32 with one max rescale and writes BF16. With pdl the launch is +//! programmatic: the grid waits on the producer (griddepcontrol.wait) before reading its scratch. +//! +//! PDL contract: griddepcontrol.wait is the kernel's first instruction, so nothing (scratch, lengths, out) is read or +//! written before every prerequisite grid has completed and flushed; a producer's early trigger +//! (griddepcontrol.launch_dependents) only lets this grid launch and wait sooner. The kernel never triggers its own +//! dependents early: the next kernel's post-wait code sees all of out, and its pre-wait prologue must not read out. +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream); + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu new file mode 100644 index 000000000000..cad2b93558ab --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu @@ -0,0 +1,1092 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// SM107 (Rubin) copy of minimaxM3/mcAllreduceNorm.cu for the sm_107 TP2 target. Two changes in the arithmetic: the +// rank sum rounds to BF16 before the residual add, the rounding points of the stock all-reduce + fused residual +// RMSNorm on this target (bf16 all-reduce output, then bf16(residual + sum)); and every call but the MoE finalize runs +// the row kernel below, whose RMS statistic is the stock fused norm's reduction tree (bitwise the stock outputs). +// +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions): they are compiled for every sm_10x target of the build, the sm_100f family target +// included (which runs on SM107). A device pass for any other architecture compiles them to a trap, never to a silent +// empty body, and the host entry point checks that the kernel binary it launches holds sm_10x code (see +// checkBinary). +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace cg = cooperative_groups; + +namespace +{ + +constexpr int kEmpty = static_cast(0x80000000u); +constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k + +// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], +// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c +// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp +// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its +// arithmetic is the same source. +constexpr int kCapThreads = 256; +constexpr int kCapIters = 3; +constexpr int kCapWarps = kCapThreads / 32; +constexpr int kCapHidden = kCapThreads * kCapIters * 8; + +// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). +__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) +{ + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int q = 0; q < 8; ++q) + { + float const v = __bfloat162float(e[q]); + ss += v * v; + } + return ss; +} + +__device__ __forceinline__ float cap_warp_sum(float v) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. +__device__ __forceinline__ void cap_store( + float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) +{ + uint4 const wr = *reinterpret_cast(w); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int q = 0; q < 8; ++q) + oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); + *reinterpret_cast(cap) = o; +} + +__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) +{ + if ((w & 0xFFFFu) == 0x8000u) + w &= 0xFFFF0000u; + if ((w >> 16) == 0x8000u) + w &= 0x0000FFFFu; + return w; +} + +__device__ __forceinline__ float bf16_lo(uint32_t w) +{ + return __uint_as_float(w << 16); +} + +// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's +// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals +// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the +// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. +__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) +{ + float value; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); + uint32_t result; + asm("{\n" + " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" + " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" + " setp.le.f32 p_zero, %1, 0f00000000;\n" + " mov.b32 bits, %1;\n" + " shr.b32 exp_biased, bits, 23;\n" + " and.b32 exp_biased, exp_biased, 255;\n" + " and.b32 mantissa, bits, 0x7FFFFF;\n" + " setp.ne.u32 p_has_mant, mantissa, 0;\n" + " selp.u32 bump, 1, 0, p_has_mant;\n" + " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" + " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" + " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" + " @p_tiny_sub mov.u32 bump, 0;\n" + " add.u32 r, exp_biased, bump;\n" + " setp.gt.u32 p_ovf, r, 254;\n" + " selp.u32 r, 254, r, p_ovf;\n" + " selp.u32 %0, 0, r, p_zero;\n" + "}\n" + : "=r"(result) + : "f"(value)); + return result; +} + +__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) +{ + float inv; + asm("{\n" + " .reg .s32 new_exp;\n" + " .reg .b32 float_bits;\n" + " .reg .pred p_zero;\n" + " setp.eq.u32 p_zero, %1, 0;\n" + " sub.s32 new_exp, 254, %1;\n" + " max.s32 new_exp, new_exp, 0;\n" + " shl.b32 float_bits, new_exp, 23;\n" + " mov.b32 %0, float_bits;\n" + " @p_zero mov.b32 %0, 0;\n" + "}\n" + : "=f"(inv) + : "r"(ue8m0)); + return inv; +} + +// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. +__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) +{ + uint16_t pair; + asm("{\n" + " .reg .f32 f0, f1;\n" + " mul.f32 f0, %1, %3;\n" + " mul.f32 f1, %2, %3;\n" + " min.f32 f0, f0, 0f43E00000;\n" + " max.f32 f0, f0, 0fC3E00000;\n" + " min.f32 f1, f1, 0f43E00000;\n" + " max.f32 f1, f1, 0fC3E00000;\n" + " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" + "}\n" + : "=h"(pair) + : "f"(lo), "f"(hi), "f"(inv)); + return pair; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ float bf16_hi(uint32_t w) +{ + return __uint_as_float(w & 0xFFFF0000u); +} + +__device__ __forceinline__ int4 ld_volatile_v4(int const* p) +{ + int4 v; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +struct FinalizeArgs +{ + __nv_bfloat16 const* fc2; // [permuted rows, H] routed-expert FC2 output (null: plain input x) + int const* perm; // [M, topk] expanded index -> permuted row, -1 = dropped + void const* scales; // [M, topk] routing weights (fp32 or bf16) + int scale_fp32; + int topk; +}; + +// grid = min(M, MAXT) * CLUSTER CTAs, cluster (CLUSTER, 1, 1), THREADS threads. +// OPTS (each bit-identical): 1 residual + norm weight loaded at entry (before the poll); 2 cluster RMS exchange by +// st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated right after the poll; 8 max-smem +// carveout so the next kernel's CTAs co-reside early. +// OPTS 16 (NOT bit-identical with the others; a different boundary): the Eagle3 drafter's boundary, whose stock is a +// plain BF16 all-reduce, then FlashInfer fused_add_rmsnorm: the RMS statistic and the normalization take the +// unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded to BF16). +// CAP: the Eagle3 capture epilogue (cap_out) is compiled in only where a call passes cap_out, so every other call +// runs the kernel without the capture branches, shared arrays or parameters in its body. +template +__global__ void __launch_bounds__(THREADS) mc_allreduce_norm_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, FinalizeArgs fin, int M, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int world, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ alignas(16) float cluster_sq[CLUSTER]; + extern __shared__ __align__(16) uint4 cap_slice[]; // [cols / 8] this CTA's updated chunks (cap_out only) + __shared__ alignas(8) unsigned long long xbar; + cg::cluster_group cluster = cg::this_cluster(); + if constexpr ((OPTS & 2) != 0) + { + if (threadIdx.x == 0) + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(b) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + } + int const crank = static_cast(cluster.block_rank()); + int const t = blockIdx.x / CLUSTER; + int const cols = H / CLUSTER; + int const c0 = crank * cols + threadIdx.x * 8; + bool const active = t < M && threadIdx.x * 8 < cols; + int64_t const half_words = static_cast(MAXT) * world * (H / 2); + + // prepushed 1: the producer (m3_mx_gemv) already pushed this rank's rows into the exchange slots; 2: and it + // launched this grid only after its pushes (its grid wait passed), so the slots are polled without a grid wait; + // 3: as 2, with x in set A and x2 in set B. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const parity = *reinterpret_cast(flags + blockIdx.x) & 1; + uint4 rraw = make_uint4(0u, 0u, 0u, 0u); + uint4 wraw = make_uint4(0u, 0u, 0u, 0u); + if constexpr ((OPTS & 1) != 0) + { + // the residual row was written grids earlier (prepushed >= 2: the producer launched this grid only after its + // own grid wait; otherwise the grid wait above) + if (active) + { + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + wraw = *reinterpret_cast(norm_w + c0); + } + } + // Push this rank's slice into every rank's buffer. + if (active && !prepushed) + { + uint4 raw; + if (fin.fc2 != nullptr) + { + // MoE finalize: fp32 sum of the routing-weighted top-k expert rows plus the shared expert, rounded once. + float f[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + if (fin.topk <= kFinMaxTopk) + { + // All top-k indices and scales, then every expert-row load in flight at once, then the same in-order + // FP32 accumulation as the loop below (one memory round trip instead of topk; bit-identical). + int prow[kFinMaxTopk]; + float sc[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + prow[k] = -1; + sc[k] = 0.f; + if (k < fin.topk) + { + int const e = t * fin.topk + k; + prow[k] = fin.perm[e]; + sc[k] = fin.scale_fp32 ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + } + } + uint4 rv[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] >= 0) + rv[k] = *reinterpret_cast(fin.fc2 + static_cast(prow[k]) * H + c0); + } +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] < 0) + continue; + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv[k]); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc[k]; + f[2 * q + 1] += v.y * sc[k]; + } + } + } + else + for (int k = 0; k < fin.topk; ++k) + { + int const e = t * fin.topk + k; + int const prow = fin.perm[e]; + if (prow < 0) + continue; + float const sc = fin.scale_fp32 + ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + uint4 const rv = *reinterpret_cast(fin.fc2 + static_cast(prow) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc; + f[2 * q + 1] += v.y * sc; + } + } + if (x2 != nullptr) + { + // As stock: the routed output is rounded to BF16 first, then the shared expert is added (one more + // BF16 rounding of the FP32 sum). + uint4 const rv = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const routed = __bfloat1622float2(__floats2bfloat162_rn(f[2 * q], f[2 * q + 1])); + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] = routed.x + v.x; + f[2 * q + 1] = routed.y + v.y; + } + } + __nv_bfloat162* o = reinterpret_cast<__nv_bfloat162*>(&raw); +#pragma unroll + for (int q = 0; q < 4; ++q) + o[q] = __floats2bfloat162_rn(f[2 * q], f[2 * q + 1]); + } + else + { + raw = *reinterpret_cast(x + static_cast(t) * H + c0); + } + if (fin.fc2 == nullptr && x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * world + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float s[8] = {}; + float sq = 0.f; + if (active) + { + // The residual row (written grids earlier) and every rank's slot are loaded together; slots still empty are + // re-polled. Ranks are then summed in rank order, as before. + if constexpr ((OPTS & 1) == 0) + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + constexpr int kMaxWorld = 8; + int* const p0 = ws_uc + parity * half_words + static_cast(t) * world * (H / 2) + c0 / 2; + int4 v[kMaxWorld]; + int4 v2[kMaxWorld]; + bool const two = prepushed == 3; + int* const p1 = p0 + 2 * half_words; // set B + uint32_t pending = (1u << world) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (r < world) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr ((OPTS & 16) != 0) + { + float const x0 = s[2 * q] + rf.x, x1 = s[2 * q + 1] + rf.y; + s[2 * q] = x0; + s[2 * q + 1] = x1; + sq += x0 * x0 + x1 * x1; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + sq += uf.x * uf.x + uf.y * uf.y; + } + } + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr ((OPTS & 16) == 0) + { + if constexpr (CAP) + cap_slice[threadIdx.x] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + if constexpr ((OPTS & 4) != 0) + { + // every CTA (active or not) rotates its call parity once its own poll is done (read only by later grids) + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffffu, sq, off); + if ((threadIdx.x & 31) == 0) + warp_sq[threadIdx.x / 32] = sq; + __syncthreads(); + if constexpr ((OPTS & 2) != 0) + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); // peers' xbar initialized (long since) + if (threadIdx.x == 0) + { + float total = 0.f; +#pragma unroll + for (int w = 0; w < THREADS / 32; ++w) + total += warp_sq[w]; + if constexpr ((OPTS & 2) == 0) + for (int c = 0; c < CLUSTER; ++c) + *cluster.map_shared_rank(&cluster_sq[crank], c) = total; + else + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t const d = static_cast(__cvta_generic_to_shared(&cluster_sq[crank])); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * 4) + : "memory"); + for (int c = 0; c < CLUSTER; ++c) + { + uint32_t rd, rb; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), + "r"(__float_as_uint(total)), "r"(rb) + : "memory"); + } + } + } + if constexpr ((OPTS & 2) == 0) + cluster.sync(); + else + { + // every peer's sum landed (st.async complete_tx on this CTA's barrier); no cluster-wide release-arrive + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t done = 0; + while (!done) + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(b) + : "memory"); + } + if constexpr ((OPTS & 16) == 0) + { + if constexpr (CAP) + { + // The row's chunks are spread over the cluster (CTA c holds chunks [c * per, (c + 1) * per)); every CTA + // forms capture_norm's statistic itself, its 128 threads standing in for capture threads ct and ct + 128. + __shared__ float cap_red[kCapWarps]; + int const per = cols / 8; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + float ca = 0.f, cb = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + { + int const ka = it * kCapThreads + static_cast(threadIdx.x); + int const kb = ka + kCapThreads / 2; + ca = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[ka % per], ka / per), ca); + cb = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[kb % per], kb / per), cb); + } + // Every peer read is done before any CTA moves on (and later exits). + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + ca = cap_warp_sum(ca); + cb = cap_warp_sum(cb); + if ((threadIdx.x & 31) == 0) + { + cap_red[threadIdx.x / 32] = ca; + cap_red[threadIdx.x / 32 + kCapWarps / 2] = cb; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + if (active) + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + if (active) + { + float total = 0.f; +#pragma unroll + for (int c = 0; c < CLUSTER; ++c) + total += cluster_sq[c]; + float const inv = rsqrtf(total / static_cast(H) + eps); + if constexpr ((OPTS & 1) == 0) + wraw = *reinterpret_cast(norm_w + c0); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + } + if (out_q != nullptr) + { + // MXFP8 copy of the normed row for the next GEMM, as the CuTe DSL mxfp8 quantize of `normed` (the tree's + // decode quantize, MX-TAC; see mxfp8_ue8m0): four threads hold one 32-column scale block. + float amax = 0.f; + if (active) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + if (active) + { + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((threadIdx.x & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } + } + if constexpr ((OPTS & 4) == 0) + { + __syncthreads(); + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +// Refuses (raises) to launch a kernel whose loaded binary holds no sm_10x code (PTX and SASS version >= 100): the +// device passes for other architectures compile the kernels to a trap, and a build that left out every sm_10x target +// must not reach the GPU at all. +template +bool checkBinary(Kernel kernel, char const* name) +{ + cudaFuncAttributes attr{}; + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, reinterpret_cast(kernel))); + TLLM_CHECK_WITH_INFO(attr.ptxVersion >= 100 && attr.binaryVersion >= 100, + "%s: the loaded binary holds no sm_10x code (PTX %d, SASS %d): rebuild with an sm_100f / sm_107a target", name, + attr.ptxVersion, attr.binaryVersion); + return true; +} + +template +void launch(MinimaxM3Sm107Tp2McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) +{ + int const M = p.numTokens, H = p.hidden; + TLLM_CHECK_WITH_INFO(H % (CLUSTER * 8) == 0 && H / CLUSTER / 8 <= THREADS, "unsupported hidden size %d", H); + TLLM_CHECK_WITH_INFO(M >= 1 && M <= MAXT, "%d tokens do not fit the %d-token bucket", M, MAXT); + auto kernel = p.capOut != nullptr ? mc_allreduce_norm_kernel + : mc_allreduce_norm_kernel; + static bool const binaryOk + = checkBinary(mc_allreduce_norm_kernel, "mc_allreduce_norm_kernel") + && checkBinary( + mc_allreduce_norm_kernel, "mc_allreduce_norm_kernel (capture)"); + (void) binaryOk; + // Both instantiations (with and without the capture epilogue) get the same attributes, whichever runs first. + static bool const configured = []() + { + auto setup = [](auto k) + { + bool ok = CLUSTER <= 8 + || cudaFuncSetAttribute(k, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess; + if ((OPTS & 8) != 0) + ok = ok + && cudaFuncSetAttribute( + k, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared) + == cudaSuccess; + return ok; + }; + bool const ok = setup(mc_allreduce_norm_kernel) + && setup(mc_allreduce_norm_kernel); + if (!ok) + cudaGetLastError(); + return ok; + }(); + TLLM_CHECK_WITH_INFO(configured, "cluster size not supported"); + FinalizeArgs const fin{p.fc2, p.perm, p.scales, p.scaleFp32 ? 1 : 0, p.topk}; + cudaLaunchConfig_t cfg = {}; + // The grid is sized to the call's rows: min(M, MAXT) clusters. Each (row, slice) slot keeps its own parity flag, + // read by the producers before they push and rotated only by the slot's CTA after it polled: a slot that a call + // does not use keeps its parity, and every rank runs the same calls, so the parities stay consistent when M + // changes between calls. + cfg.gridDim = dim3((M < MAXT ? M : MAXT) * CLUSTER); + cfg.blockDim = dim3(THREADS); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO((OPTS & 16) == 0 && H == kCapHidden && 2 * THREADS == kCapThreads, + "the capture epilogue takes H %d on the 128-thread cluster kernel, not the drafter boundary", kCapHidden); + cfg.dynamicSmemBytes = static_cast(H / CLUSTER / 8) * sizeof(uint4); + } + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = CLUSTER; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, fin, M, H, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, p.world, p.outNormed, p.outResidual, p.outQuant, p.outSf, (H / 32 + 3) / 4, + p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +// ws_uc / ws_mc are this rank's unicast and multicast views of ONE exchange buffer (they alias): no __restrict__. +// prepushed >= 2 contract: the producer (m3_mx_gemv's o_proj push, the m3_moe engine's finalize push) passes its own +// grid dependency, pushes its rows into the exchange, and only then releases this grid (griddepcontrol +// .launch_dependents); this kernel therefore skips griddepcontrol.wait and reads only the pushed (Lamport-polled) +// words, its parity flags and the residual / norm weight, which were written two or more grids earlier. +// +// The TP2 row kernel: one CTA of H / 8 = 768 threads per token row, used for every call but the MoE finalize. Its +// residual add + RMSNorm is TRT-LLM's rms_norm_kernel (customAllReduceKernels.cu), the stock fused norm after the +// all-reduce on this target, operation for operation at H 6144: thread j owns columns [8j, 8j + 8), its sum of squares +// is acc += v * v over them in order, then the warp butterfly 16..1, the 24 warp sums through shared memory into a +// second 32-lane butterfly (lanes >= 24 zero), rsqrtf(acc / H + eps) and (v * denom) * w. With the BF16 rank sum +// before the residual add, the normed rows, the updated residual and the MXFP8 copy are bitwise the stock ones. The +// exchange slots and the per-(row, slice) parity flags keep the cluster kernel's layout (slice = 96 threads), so +// producers that push (the o_proj GEMV) are unchanged. +constexpr int kSlices = 8; +constexpr int kSliceThreads = 96; +constexpr int kRowThreads = 768; + +// grid = M CTAs (one per row), THREADS = H / 8 threads. WORLD ranks (compile time). DRAFTER (opts bit 16): the Eagle3 +// drafter boundary, whose stock is a plain BF16 all-reduce then FlashInfer fused_add_rmsnorm: the RMS statistic and +// the normalization take the unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded). +// ws_uc / ws_mc alias (one exchange buffer): no __restrict__. +template +__global__ void __launch_bounds__(THREADS, 1) mc_allreduce_norm_row_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int maxt, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ float total_sh; + extern __shared__ __align__(16) uint4 cap_row[]; // [THREADS] the updated row's chunks (cap_out only) + int const t = blockIdx.x; + int const j = threadIdx.x; + int const c0 = j * 8; + int const slice = j / kSliceThreads; + int64_t const half_words = static_cast(maxt) * WORLD * (H / 2); + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); + int const parity = *reinterpret_cast(flags + t * kSlices + slice) & 1; + uint4 const rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + uint4 const wraw = *reinterpret_cast(norm_w + c0); + if (!prepushed) + { + uint4 raw = *reinterpret_cast(x + static_cast(t) * H + c0); + if (x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * WORLD + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } + asm volatile("griddepcontrol.launch_dependents;"); + float s[8] = {}; + float sq = 0.f; + { + int* const p0 = ws_uc + parity * half_words + static_cast(t) * WORLD * (H / 2) + c0 / 2; + int* const p1 = p0 + 2 * half_words; // set B + bool const two = prepushed == 3; + int4 v[WORLD]; + int4 v2[WORLD]; + uint32_t pending = (1u << WORLD) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr (DRAFTER) + { + s[2 * q] += rf.x; + s[2 * q + 1] += rf.y; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + } + } +#pragma unroll + for (int e = 0; e < 8; ++e) + sq += s[e] * s[e]; + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr (!DRAFTER) + { + if constexpr (CAP) + cap_row[j] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + // TRT-LLM reduce_fusion::block_reduce_sum. + sq += __shfl_xor_sync(~0, sq, 16); + sq += __shfl_xor_sync(~0, sq, 8); + sq += __shfl_xor_sync(~0, sq, 4); + sq += __shfl_xor_sync(~0, sq, 2); + sq += __shfl_xor_sync(~0, sq, 1); + if ((j & 31) == 0) + warp_sq[j / 32] = sq; + __syncthreads(); + { + float val = (j & 31) < THREADS / 32 ? warp_sq[j & 31] : 0.f; + val += __shfl_xor_sync(~0, val, 16); + val += __shfl_xor_sync(~0, val, 8); + val += __shfl_xor_sync(~0, val, 4); + val += __shfl_xor_sync(~0, val, 2); + val += __shfl_xor_sync(~0, val, 1); + if (j == 0) + total_sh = val; + } + // The parity of this row's slices rotates once every thread has read the exchange (later grids only). + if (j % kSliceThreads == 0) + flags[t * kSlices + slice] = parity ^ 1; + __syncthreads(); + if constexpr (!DRAFTER) + { + if constexpr (CAP) + { + // The whole row is in this CTA: threads 0..255 run capture_norm's chains over the staged chunks. + __shared__ float cap_red[kCapWarps]; + if (j < kCapThreads) + { + float css = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + css = cap_chunk_sq(cap_row[it * kCapThreads + j], css); + css = cap_warp_sum(css); + if ((j & 31) == 0) + cap_red[j / 32] = css; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + float const inv = rsqrtf(total_sh / H + eps); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + if (out_q != nullptr) + { + float amax = 0.f; +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((j & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +void launchRow( + MinimaxM3Sm107Tp2McAllreduceNormParams const& p, int32_t maxTokens, bool drafter, bool pdl, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(p.hidden == kRowThreads * 8 && p.world == 2, "the row kernel takes H 6144 at 2 ranks"); + static bool const binaryOk + = checkBinary(mc_allreduce_norm_row_kernel<2, kRowThreads, false, false>, "mc_allreduce_norm_row_kernel") + && checkBinary( + mc_allreduce_norm_row_kernel<2, kRowThreads, true, false>, "mc_allreduce_norm_row_kernel (drafter)") + && checkBinary( + mc_allreduce_norm_row_kernel<2, kRowThreads, false, true>, "mc_allreduce_norm_row_kernel (capture)"); + (void) binaryOk; + TLLM_CHECK_WITH_INFO(p.numTokens >= 1 && p.numTokens <= maxTokens, "%d tokens do not fit the %d-token bucket", + p.numTokens, maxTokens); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(p.numTokens); + cfg.blockDim = dim3(kRowThreads); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO( + !drafter && p.hidden == kCapHidden, "the capture epilogue takes H %d, not the drafter", kCapHidden); + cfg.dynamicSmemBytes = static_cast(kRowThreads) * sizeof(uint4); + } + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + // The capture epilogue never runs on the drafter boundary (checked above), so there are three instantiations. + auto kernel = drafter ? mc_allreduce_norm_row_kernel<2, kRowThreads, true, false> + : p.capOut != nullptr ? mc_allreduce_norm_row_kernel<2, kRowThreads, false, true> + : mc_allreduce_norm_row_kernel<2, kRowThreads, false, false>; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, p.hidden, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, maxTokens, p.outNormed, p.outResidual, p.outQuant, p.outSf, + (p.hidden / 32 + 3) / 4, p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +constexpr int kCluster = 8; +constexpr int kThreads = 128; + +} // namespace + +bool isMinimaxM3Sm107Tp2McAllreduceNormSupported(int32_t maxTokens, int32_t opts) +{ + return (maxTokens == 8 && (opts == 7 || opts == 15 || opts == 23 || opts == 31)) || (maxTokens == 16 && opts == 15) + || (maxTokens == 48 && opts == 15); +} + +void invokeMinimaxM3Sm107Tp2McAllreduceNorm(MinimaxM3Sm107Tp2McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream) +{ + if (params.fc2 == nullptr && params.world == 2 && params.hidden == kRowThreads * 8) + { + launchRow(params, maxTokens, (opts & 16) != 0, pdl, stream); + return; + } + if (maxTokens == 8 && opts == 7) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 23) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 31) + { + launch(params, pdl, stream); + } + else if (maxTokens == 16 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 48 && opts == 15) + { + launch(params, pdl, stream); + } + else + { + TLLM_THROW("unsupported (bucket %d, opts %d)", maxTokens, opts); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h new file mode 100644 index 000000000000..e8ade9c62c3e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h @@ -0,0 +1,83 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Operands of the MiniMax-M3 multicast all-reduce + residual + RMSNorm kernel (decode, at most 48 tokens). +//! +//! Per rank the kernel forms the partial `bf16(x + x2)` (or, in finalize mode, +//! `bf16(bf16(sum_k scale * fc2[perm]) + x2)`, the routed sum rounded before the shared expert is added, as stock), +//! pushes it through the multicast mapping `wsMc` into every rank's exchange, polls its own unicast view `wsUc`, +//! sums the ranks in FP32 in rank order, adds the residual (BF16 round) and applies RMSNorm. With `outQuant` it +//! also writes the MXFP8 copy of the normed rows (E4M3 [numTokens, hidden] and swizzled UE8M0 scales). +struct MinimaxM3Sm107Tp2McAllreduceNormParams +{ + __nv_bfloat16 const* x{nullptr}; //!< [numTokens, hidden] partial; in finalize mode the FC2 rows [rows, hidden]. + __nv_bfloat16 const* x2{nullptr}; //!< Optional second addend [numTokens, hidden]. + __nv_bfloat16 const* fc2{nullptr}; //!< Finalize mode: the permuted FC2 output (same buffer as x); null otherwise. + int32_t const* perm{nullptr}; //!< Finalize mode: [numTokens, topk] expanded index -> permuted row, -1 dropped. + void const* scales{nullptr}; //!< Finalize mode: [numTokens, topk] routing weights (fp32 or bf16). + bool scaleFp32{false}; + int32_t topk{0}; + int32_t numTokens{0}; + int32_t hidden{0}; + __nv_bfloat16 const* residual{nullptr}; //!< [numTokens, hidden] + __nv_bfloat16 const* normWeight{nullptr}; //!< [hidden] + float eps{0.F}; + int32_t* wsUc{nullptr}; //!< This rank's unicast view of the exchange (int32 words, re-armed after reading). + int32_t* wsMc{nullptr}; //!< The multicast mapping of the exchange. + int32_t* flags{nullptr}; //!< Per-CTA call-parity counters, [maxTokens * 8]. + int32_t rank{0}; + int32_t world{0}; + __nv_bfloat16* outNormed{nullptr}; + __nv_bfloat16* outResidual{nullptr}; + uint8_t* outQuant{nullptr}; //!< Optional E4M3 [numTokens, hidden]. + uint8_t* outSf{nullptr}; //!< Optional UE8M0 scales, 128x4-swizzled. + //! 0: push x (+ x2) here; 2: a producer already pushed x and launched this grid after its pushes; 3: as 2, with x + //! in exchange set A and x2 in set B. + int32_t prepushed{0}; + //! Optional Eagle3 capture of a capture layer's boundary: capOut[t * capStride + c] = the m3_spec_tail + //! capture_norm of the updated residual row with weight capWeight [hidden] and capEps, bit for bit (hidden 6144, + //! not the drafter boundary). + __nv_bfloat16* capOut{nullptr}; + int64_t capStride{0}; + __nv_bfloat16 const* capWeight{nullptr}; + float capEps{0.F}; +}; + +//! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). +bool isMinimaxM3Sm107Tp2McAllreduceNormSupported(int32_t maxTokens, int32_t opts); + +//! Launch with 8-CTA clusters of 128 threads, one cluster per token. `maxTokens` selects the exchange bucket; +//! `opts` the epilogue variant (every variant is bit-identical). +void invokeMinimaxM3Sm107Tp2McAllreduceNorm(MinimaxM3Sm107Tp2McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu new file mode 100644 index 000000000000..da30d840cc2e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu @@ -0,0 +1,339 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h" + +#include +#include +#include + +// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP2 heads. +// +// The SM107 TP2 target's copy of minimaxM3/producerPw.cu (the sm_103 TP4 / TP8 targets keep theirs). Same math and +// outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with --use_fast_math like +// the rest of the library), specialized to 32 query, 2 KV and 2 index heads per rank (one index-K head), head_dim 128, +// rotary_dim 64, page 128, NVFP4 main KV and FP8 index-K. One warp per (token, packed head); K and V rows go to their +// head's block of the HND pools (the native kernel's head strides). +// +// The decode producer is latency bound: after griddepcontrol.wait the native kernel runs a chain of dependent global +// loads (packed row -> norm weight -> position id -> RoPE row -> cache slot -> quant scale). This kernel reads the +// constant norm weights and KV quant scale before griddepcontrol.wait and issues the packed row, position id and cache +// slot loads together after it, followed by the RoPE row; with `prewaitStepInputs` the position id, the slot and the +// RoPE row are read before the wait too, leaving only the packed row after it. With `early` it also calls +// griddepcontrol.launch_dependents on entry: consumers read its outputs only after their own griddepcontrol.wait, which +// waits for this grid to complete, so the early trigger only lets the next kernel launch and run its pre-wait prologue +// while this one runs. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace producerPw2 +{ +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotaryDim = 64; +constexpr int kPageSize = 128; +constexpr int kElems = kHeadDim / 32; +constexpr int kBlock = 256; + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ uint32_t sf_group_mask4() +{ + uint32_t laneId; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return 0xFu << (laneId & ~3u); +} + +__device__ __forceinline__ uint32_t fp32x8_to_e2m1(float (&a)[8]) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(a[0]), "f"(a[1]), "f"(a[2]), "f"(a[3]), "f"(a[4]), "f"(a[5]), "f"(a[6]), "f"(a[7])); + return val; +#else + return 0; +#endif +} + +__device__ __forceinline__ void store_fp8x4(__nv_fp8_e4m3* out, int64_t off, float const (&e)[kElems]) +{ + __nv_fp8x2_e4m3 const low(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const high(make_float2(e[2], e[3])); + *reinterpret_cast(out + off) = static_cast(low.__x) | (static_cast(high.__x) << 16); +} + +__device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(e[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = sf_group_mask4(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + *scaleOut = encodedScale.__x; + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float q[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + q[i] = __bfloat162float(rounded[i]) * outputScale; + return static_cast(fp32x8_to_e2m1(q)); +} + +template +__global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (kEarly) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * (kBlock / 32) + warpId; + int const tokenIdx = globalWarp / kTotalHeads; + int const localHead = globalWarp % kTotalHeads; + if (tokenIdx >= p.numTokens) + return; + + constexpr int kBegin = kHeadsQ; + constexpr int vBegin = kBegin + kHeadsKV; + constexpr int indexQBegin = vBegin + kHeadsKV; + constexpr int indexKHead = indexQBegin + kHeadsIndex; + bool const isQ = localHead < kBegin; + bool const isK = localHead >= kBegin && localHead < vBegin; + bool const isV = localHead >= vBegin && localHead < indexQBegin; + bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; + bool const isIndexK = localHead == indexKHead; + + // Constant data read before the wait: norm weights and the KV quant scale. + float w[kElems] = {}; + float quantScale = 0.0F; + __nv_bfloat16 const* normWeight + = isQ ? p.qWeight : (isK ? p.kWeight : (isIndexQ ? p.indexQWeight : p.indexKWeight)); + if (!isV) + { + uint2 const wv = *reinterpret_cast(normWeight + laneId * kElems); + float2 const w01 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.x)); + float2 const w23 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.y)); + w[0] = w01.x; + w[1] = w01.y; + w[2] = w23.x; + w[3] = w23.y; + } + if (isK || isV) + { + quantScale = p.kvQuantScale[isV ? 2 : 1]; + } + + int64_t const inputOffset = (static_cast(tokenIdx) * kTotalHeads + localHead) * kHeadDim + laneId * kElems; + bool const rotates = !isV && laneId * kElems < kRotaryDim; + bool const slotted = isK || isV || isIndexK; + int const coefficient = (laneId * kElems) % (kRotaryDim / 2); + // kPrewait: the position id, the cache slot and the RoPE row are read before the wait. They are per-step metadata + // written before the replay (host-to-device) or by plain-launched kernels earlier in the step (position_ids by the + // step's input preprocessing, the slots by msa_kv_update or its torch fallback), so they are final before this + // grid can start whatever its predecessors trigger; the RoPE table is constant. + int positionId = 0; + int slot = 0; + float4 c = make_float4(0.0F, 0.0F, 0.0F, 0.0F); + float4 s = c; + auto const loadStepInputs = [&]() + { + positionId = rotates ? p.positionIds[tokenIdx] : 0; + slot = slotted ? p.outCacheLoc[tokenIdx] : 0; + if (rotates) + { + int64_t const ropeRow = static_cast(positionId) * kRotaryDim; + c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); + s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); + } + }; + if constexpr (kPrewait) + { + loadStepInputs(); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + // The packed row comes from the predecessor grid; without kPrewait the step inputs are issued with it, so the + // native kernel's serial chain of dependent loads becomes two round trips. + uint2 const packed = *reinterpret_cast(p.packed + inputOffset); + if constexpr (!kPrewait) + { + loadStepInputs(); + } + float cosine[kElems] = {c.x, c.y, c.z, c.w}; + float sine[kElems] = {s.x, s.y, s.z, s.w}; + + float e[kElems]; + float sumSquares = 0.0F; +#pragma unroll + for (int pair = 0; pair < 2; ++pair) + { + float2 const v = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); + e[2 * pair] = v.x; + e[2 * pair + 1] = v.y; + if (!isV) + sumSquares += v.x * v.x + v.y * v.y; + } + + if (!isV) + { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + sumSquares += __shfl_xor_sync(0xffffffff, sumSquares, mask, 32); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kHeadDim) + p.eps); +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + constexpr int kPairOffset = (kRotaryDim / 2) / kElems; +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + int const dim = laneId * kElems + i; + float paired = __shfl_xor_sync(0xffffffff, e[i], kPairOffset); + if (dim < kRotaryDim) + { + if (dim < kRotaryDim / 2) + paired = -paired; + // The native kernel's contraction of e * cos + paired * sin; spelled out so this build matches it. + e[i] = __fmaf_rn(paired, sine[i], __fmul_rn(e[i], cosine[i])); + } + } + } + + if (isQ) + { + store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); + return; + } + if (isIndexQ) + { + store_fp8x4(p.indexQOut, + (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); + return; + } + + if (slot < 0) + return; + int const page = slot / kPageSize; + if (page >= p.numPages) + return; + int const withinPage = slot & (kPageSize - 1); + if (isIndexK) + { + store_fp8x4(p.indexK, + static_cast(page) * p.indexPageStride + static_cast(withinPage) * p.indexTokenStride + + laneId * kElems, + e); + return; + } + + int const head = isK ? localHead - kBegin : localHead - vBegin; + int const plane = isV ? 1 : 0; + int const scaleColumn = laneId / 4; + int64_t const scaleWithinPage = isV + ? static_cast(withinPage / 4) * 32 + static_cast(scaleColumn) * 4 + withinPage % 4 + : static_cast(withinPage) * 8 + scaleColumn; + int64_t const scaleOffset = static_cast(page) * p.scalePageStride + + static_cast(plane) * p.scalePlaneStride + static_cast(head) * p.scaleHeadStride + + scaleWithinPage; + uint8_t* scaleOut = laneId % 4 == 0 ? p.kvScale + scaleOffset : nullptr; + uint16_t const packedFp4 = quantize_fp4x4(e, quantScale, scaleOut); + int64_t const dataOffset = static_cast(page) * p.dataPageStride + + static_cast(plane) * p.dataPlaneStride + static_cast(head) * p.dataHeadStride + + static_cast(withinPage) * p.dataTokenStride + laneId * 2; + *reinterpret_cast(p.kvData + dataOffset) = packedFp4; +#endif +} + +template +cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) +{ + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; + int const warps = numTokens * kTotalHeads; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); + cfg.blockDim = dim3(kBlock); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); +} + +} // namespace + +cudaError_t invokeProducerPw2( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream) +{ + if (numTokens <= 0) + return cudaSuccess; + if (prewaitStepInputs) + { + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); + } + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); +} + +} // namespace producerPw2 +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h new file mode 100644 index 000000000000..9e94868af843 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP2 heads (32 query, 2 KV and +//! 2 index heads per rank) with its constant loads before the PDL grid-dependency wait and an optional early dependents +//! trigger. Same outputs, bit for bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw2.cu. +namespace producerPw2 +{ + +struct Params +{ + __nv_bfloat16 const* packed; + __nv_fp8_e4m3* qOut; + __nv_fp8_e4m3* indexQOut; + uint8_t* kvData; + uint8_t* kvScale; + __nv_fp8_e4m3* indexK; + int const* outCacheLoc; + float const* kvQuantScale; + int64_t dataPageStride, dataPlaneStride, dataHeadStride, dataTokenStride; + int64_t scalePageStride, scalePlaneStride, scaleHeadStride; + int64_t indexPageStride, indexTokenStride, numPages; + int numTokens; + float eps; + __nv_bfloat16 const* qWeight; + __nv_bfloat16 const* kWeight; + __nv_bfloat16 const* indexQWeight; + __nv_bfloat16 const* indexKWeight; + float const* rotaryCosSin; + int const* positionIds; +}; + +//! The heads per rank this producer is built for. +constexpr int kHeadsQ = 32; +constexpr int kHeadsKV = 2; +constexpr int kHeadsIndex = 2; + +//! Launch over `numTokens` tokens, one warp per (token, packed head). `prewaitStepInputs` reads the position ids, the +//! cache slots and the RoPE rows before the grid-dependency wait (they must not be written by a PDL predecessor). +cudaError_t invokeProducerPw2( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream); + +} // namespace producerPw2 +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu new file mode 100644 index 000000000000..878282b18839 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu @@ -0,0 +1,596 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h" + +#include "tensorrt_llm/common/assert.h" + +#include +#include + +// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). +// +// sm_107 (Rubin) TP2 copy of kernels/minimaxM3/sampledAccept.cu: identical algorithm and arithmetic; the vocab shard +// limit is raised to the TP2 shard (200064 / 2 = 100032 -> kMaxVs 100352), so a CTA stages up to kMaxVs / kSplits +// BF16 logits (25 KB at 8 splits, 50 KB at 4) and the launch opts in to that dynamic shared memory. +// +// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +// active must stay on the generic path; the caller decides. +// +// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens +// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the +// nucleus with probability >= top_p, so most rows decide on their second exchange. +// +// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based +// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one +// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and +// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the +// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing +// and two call halves keep a rank one call ahead from overwriting a slot still being read. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace sampledAccept +{ +namespace +{ +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kMaxRec = kMaxWorld * kMaxSplits; +constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) +constexpr int kEx = 4; // words per record: key, logit of key, f2, f3 +constexpr int kTagBits = 64 - 32 - kIdBits; +constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard +// Record polls trap after this many spins (the multicast all-reduce's bound): a missing writer becomes a CUDA error, +// never a hang. +constexpr long long kMaxSpins = 1LL << 31; + +// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. +__host__ __device__ inline long long half_words(int world) +{ + return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * kEx; +} + +__device__ __forceinline__ uint32_t ordered(float v) +{ + uint32_t const b = __float_as_uint(v); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +// Philox4x32-10. +__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) +{ +#pragma unroll + for (int i = 0; i < 10; ++i) + { + uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; + uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; + c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); + k.x += 0x9E3779B9u; + k.y += 0xBB67AE85u; + } + return c; +} + +// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. +__device__ __forceinline__ float uniform(uint32_t bits) +{ + return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); +} + +// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. +__device__ __forceinline__ unsigned long long gkey(float lg, int id) +{ + return (static_cast(ordered(lg)) << kIdBits) + | static_cast((1u << kIdBits) - 1u - static_cast(id)); +} + +__device__ __forceinline__ int key_token(unsigned long long k) +{ + return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); +} + +__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) +{ + asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); +} + +__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) +{ + unsigned long long v; + asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); + return v; +} + +template +struct Shared +{ + __nv_bfloat16 row[kMaxVs / kSplits]; + float redf[kWarps]; + unsigned long long redk[kWarps]; + unsigned long long rk[kMaxRec]; + float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; + float wsum[kWarps]; + int owner; + int pick; + float pickx; + bool last; +}; + +template +__device__ __forceinline__ float block_max(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = s.redf[0]; + for (int i = 1; i < kWarps; ++i) + r = fmaxf(r, s.redf[i]); + __syncthreads(); + return r; +} + +// Fixed-order block sum: identical result for identical inputs. +template +__device__ __forceinline__ float block_sum(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = 0.f; + for (int i = 0; i < kWarps; ++i) + r += s.redf[i]; + __syncthreads(); + return r; +} + +template +__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + if ((threadIdx.x & 31) == 0) + s.redk[threadIdx.x >> 5] = v; + __syncthreads(); + unsigned long long r = s.redk[0]; + for (int i = 1; i < kWarps; ++i) + r = s.redk[i] > r ? s.redk[i] : r; + __syncthreads(); + return r; +} + +// Exclusive prefix of v over the threads in thread order, and the total. +template +__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + float inc = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + float const t = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) + inc += t; + } + if (lane == 31) + s.wsum[warp] = inc; + __syncthreads(); + float base = 0.f; + total = 0.f; + for (int i = 0; i < kWarps; ++i) + { + base += i < warp ? s.wsum[i] : 0.f; + total += s.wsum[i]; + } + __syncthreads(); + return base + inc - v; +} + +// A row's kSplits CTAs wait for each other's records, so they launch as one thread-block cluster: the hardware +// schedules a cluster's CTAs together, whatever other streams hold. Rows never wait on other rows. +template +__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) +{ + extern __shared__ __align__(16) unsigned char smem[]; + Shared& s = *reinterpret_cast*>(smem); + float constexpr kNegInf = -__builtin_huge_valf(); + int const tid = threadIdx.x; + int const row = blockIdx.x / kSplits; + int const split = blockIdx.x % kSplits; + int const world = a.world; + int const nrec = world * kSplits; + // The row's sampling parameters are step inputs (the spec metadata, written before the target forward), not the + // predecessor head GEMM's output, so with prewait_params they load before the dependency wait and overlap the + // head's tail; the same values either way. + float temperature = 0.f, topP = 0.f; + long long seed_in = 0, offset_in = 0; + auto load_params = [&]() + { + temperature = a.temperatures[row]; + topP = a.top_ps[row]; + seed_in = a.seeds[row]; + offset_in = a.offsets[row]; + }; + if (a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (!a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // early_trigger: the dependent grid (the first draft pass's input) may launch now and park at its own + // griddepcontrol.wait, which still returns only once this grid has completed and flushed; it skips the launch + // latency after this kernel ends. Otherwise dependents are released at the acceptance epilogue (the CTAs of + // splits 1.. exit without releasing). + if (a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const tempInv = 1.0f / (temperature + kTempEpsilon); + unsigned long long const seed = static_cast(seed_in); + unsigned long long const offset = static_cast(offset_in); + unsigned int const call = *reinterpret_cast(a.counters + 1); + // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of + // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round + // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten + // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. + auto ktag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; + }; + auto ftag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; + }; + long long const rec_stride = static_cast(world) * kMaxSplits * kEx; + unsigned long long* const xu + = a.ws_uc + static_cast(call & 1u) * half_words(world) + static_cast(row) * rec_stride; + unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * kEx; + long long const round_stride = static_cast(kMaxRows) * rec_stride; + + // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- + __nv_bfloat16 const* src = a.logits + row * a.ld; + int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; + int const s0 = min(split * sl_len, a.Vs); + int const n = min(s0 + sl_len, a.Vs) - s0; + int const id0 = a.vocab_offset + s0; + float tmax = kNegInf; + bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; + int const nvec = vec ? n / 8 : 0; + for (int v = tid; v < nvec; v += kThreads) + { + uint4 const w = *reinterpret_cast(src + s0 + 8 * v); + *reinterpret_cast(&s.row[8 * v]) = w; + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); +#pragma unroll + for (int j = 0; j < 8; ++j) + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); + } + for (int i = nvec * 8 + tid; i < n; i += kThreads) + { + s.row[i] = src[s0 + i]; + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); + } + float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier + float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z + unsigned long long ckey = 0ULL; + int const gid = a.rank * kMaxSplits + split; + for (int r = 0;; ++r) + { + // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) + float const ref = r == 0 ? smax : M; + float const wmin = r == 0 ? 0.f : wc; + float ts = 0.f; + for (int i = tid; i < n; i += kThreads) + { + float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); + ts += w > wmin ? w : 0.f; + } + float total; + float const excl = block_excl_scan(ts, total, s); + uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), + static_cast(offset), static_cast(offset >> 32)), + make_uint2(static_cast(seed), static_cast(seed >> 32))); + float const tgt = uniform(rnd.x) * total; + if (tid == 0) + s.owner = -1; + __syncthreads(); + if (ts > 0.f && excl <= tgt) + atomicMax(&s.owner, tid); + __syncthreads(); + if (tid == s.owner) + { + float acc = excl; + int pick = -1; + float px = 0.f; + if constexpr (kIlpWalk) + { + // SA-PW: the owner's strided tokens 8 at a time. The 8 x and w (the same expressions) are + // independent, so their MUFU work overlaps; the serial sum then adds them in the same order with + // the same early exit: the same acc sequence, pick and px as the one-token walk, bit for bit. + constexpr int kW = 8; + bool done = false; + for (int i0 = tid; i0 < n && !done; i0 += kW * kThreads) + { + float xs[kW], wv[kW]; +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + xs[k] = i < n ? __fmul_rn(__bfloat162float(s.row[i]), tempInv) : 0.f; + wv[k] = __expf(xs[k] - ref); + } +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + if (!done && i < n && wv[k] > wmin) + { + acc += wv[k]; + pick = i; + px = xs[k]; + if (acc > tgt) + done = true; + } + } + } + } + else + { + for (int i = tid; i < n; i += kThreads) + { + float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); + float const w = __expf(x - ref); + if (w > wmin) + { + acc += w; + pick = i; + px = x; + if (acc > tgt) + break; + } + } + } + s.pick = pick; + s.pickx = px; + } + __syncthreads(); + unsigned long long sb = 0ULL; + if (total > 0.f && s.owner >= 0) + sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); + float const smab = total; + unsigned long long const ktag = ktag_of(r); + unsigned long long const ftag = ftag_of(r); + unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; + if (tid == 0) + { + unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; + st_sys(xw + 0, (sb << kTagBits) | ktag); + st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); + st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); + st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); + } + // ---- every record of the row, every rank --------------------------------------------------------------- + if (tid < nrec * kEx) + { + int const q = tid / kEx, f = tid % kEx; + unsigned long long v; + if (f == 0) + { + long long spins = 0; + do + { + v = ld_sys(xr + (q / kSplits * kMaxSplits + q % kSplits) * kEx); + if (++spins > kMaxSpins) + __trap(); + } while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); + s.rk[q] = v >> kTagBits; + } + else + { + long long spins = 0; + do + { + v = ld_sys(xr + (q / kSplits * kMaxSplits + q % kSplits) * kEx + f); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFFFFF00000000ULL) != ftag); + float const fv = __uint_as_float(static_cast(v)); + (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; + } + } + __syncthreads(); + // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- + unsigned long long gk = 0ULL; + float gv = 0.f; + for (int q = 0; q < nrec; ++q) + { + if (s.rk[q] > gk) + { + gk = s.rk[q]; + gv = s.rv[q]; + } + } + bool accept; + if (r == 0) + { + M = kNegInf; + for (int q = 0; q < nrec; ++q) + M = fmaxf(M, s.r2[q]); + float Z = 0.f; + for (int q = 0; q < nrec; ++q) + Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); + target = topP * Z; + accept = !(topP < 1.0f); // no nucleus cut: the first draw stands + ckey = gk; + cv = gv; + } + else + { + float S = 0.f; // mass strictly above the current candidate + for (int q = 0; q < nrec; ++q) + S += s.r2[q]; + accept = S <= target || gk == 0ULL; + if (!accept) + { + ckey = gk; + cv = gv; + } + } + wc = __expf(cv - M); + __syncthreads(); // records consumed before the next round's writes to s + if (accept) + { + if (split != 0) + return; + if (tid == 0) + { + a.tokens[row] = key_token(ckey); + if (a.extra_attempts != nullptr && r > 1) + atomicAdd(a.extra_attempts, r - 1); + } + break; + } + } + + // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (!a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (tid == 0) + s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; + __syncthreads(); + if (!s.last) + return; + __threadfence(); + if (tid == 0) + { + a.counters[0] = 0u; + a.counters[1] = call + 1u; + } + if (tid < a.nreq) + { + int const* tok = a.tokens; + int const b = tid; + int const k1 = a.draft_len + 1; + int nacc = 1; + bool run = true; + for (int k = 0; k < a.draft_len; ++k) + { + int const tt = *reinterpret_cast(tok + b * k1 + k); + a.acc_tokens[b * k1 + k] = tt; + run = run && a.draft[b * a.draft_len + k] == tt; + nacc += run ? 1 : 0; + } + a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); + if (a.force_vec != nullptr) + nacc = a.force_vec[b]; + else if (a.force_total > 0) + nacc = a.force_total; + a.num_acc[b] = nacc; + a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; + if (a.next_new != nullptr) + a.next_new[b * a.next_stride] + = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); + } +} + +template +cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) +{ + static int max_clusters = 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(a.rows * kSplits); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Shared); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + if (max_clusters == 0) + { + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } + } + TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); + return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); +} + +} // namespace + +long long sampledAcceptWsWords(int world) +{ + return 2 * half_words(world); +} + +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) +{ + // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); + // more rows: 4 CTAs per row. + if (a.rows <= 32) + { + return a.ilp_walk ? launch<8, true>(a, usePdl, stream) : launch<8, false>(a, usePdl, stream); + } + return a.ilp_walk ? launch<4, true>(a, usePdl, stream) : launch<4, false>(a, usePdl, stream); +} + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h new file mode 100644 index 000000000000..cd1e8ec7fefa --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h @@ -0,0 +1,82 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +//! Fused sampled (temperature + top-p) verification of the MiniMax-M3 Eagle3 target head over the TP-sharded +//! vocabulary, decode rows <= kMaxRows; see sampledAccept.cu. +namespace sampledAccept +{ + +constexpr int kMaxSplits = 8; // CTAs per verify row (at most) +constexpr int kMaxRows = 48; // verify rows per call +constexpr int kMaxWorld = 8; // TP ranks +constexpr int kMaxVs = 100352; // vocab shard width (TP2: 200064 / 2 = 100032, padded to a multiple of 512) +constexpr int kIdBits = 18; // token ids < 2^18 + +struct Args +{ + __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard + long long ld; + int rows, Vs, vocab_offset; + float const* temperatures; + float const* top_ps; + long long const* seeds; + long long const* offsets; + unsigned long long* ws_uc; + unsigned long long* ws_mc; + unsigned int* counters; // [0] finished rows, [1] call + int* tokens; // [rows] sampled token per row + int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows + int rank, world; + int const* draft; + int const* batch_idx; + int* acc_tokens; + int* num_acc; + long long* gather_ids; + int draft_len, nreq, force_total; + int* next_new; + int next_stride; + int const* force_vec; + int prewait_params; // nonzero: the row's sampling parameters load before the PDL dependency wait + int early_trigger; // nonzero: dependents are released right after the dependency wait, not at the epilogue + int ilp_walk; // nonzero: the owner walk takes its tokens 8 at a time (same weights, same summation order) +}; + +//! Words (uint64) of the multicast exchange workspace for `world` ranks: two call halves. +long long sampledAcceptWsWords(int world); + +//! Launch the sampled verification over a.rows * (8 or 4) co-resident CTAs (8 per row up to 32 rows). +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream); + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.cu new file mode 100644 index 000000000000..e708083ef9ff --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.cu @@ -0,0 +1,246 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. +// +// One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +// thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups `block % 16` is a lower +// bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared +// memory and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every +// candidate against the whole row (re-read). + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr float kInitScore = 1.0e30f; +constexpr float kLocalScore = 1.0e29f; +constexpr uint32_t kNegInfKey = 0xff800000u; +constexpr int kItems = 9; // keys per thread: rows up to 9216 blocks (1M tokens + draft tokens) + +// Descending order key: a smaller key is a larger score (NaN above +inf), as in the reference op. +__device__ __forceinline__ uint32_t orderedKey(float x) +{ + uint32_t const b = __float_as_uint(x); + return (b & 0x80000000u) ? b : (~b & 0x7fffffffu); +} + +__device__ __forceinline__ float applyForcing(float raw, int block, int initBlocks, long long localStart) +{ + if (block >= localStart) + return kLocalScore; + if (block < initBlocks) + return kInitScore; + return raw; +} + +struct SelectSmem +{ + uint32_t groupMin[16]; + unsigned long long cand[1024]; + int numCand; + unsigned long long picked[16]; +}; + +// One row with up to kItems * 1024 valid blocks; the caller picks kItems from the row's valid count. +template +__device__ __forceinline__ void selectRow(SelectSmem& sm, float const* rowBase, long long blockStride, int valid, + int initBlocks, long long localStart, bool loadedTrigger) +{ + int const tid = threadIdx.x, lane = tid & 31; + uint32_t key[kItems]; + uint32_t localMin = 0xffffffffu; +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + key[s] = block < valid ? orderedKey(applyForcing(rowBase[block * blockStride], block, initBlocks, localStart)) + : 0xffffffffu; + localMin = min(localMin, key[s]); + } + if (valid <= 16) + { + if (tid < valid) + sm.picked[tid] = (static_cast(key[0]) << 32) | static_cast(tid); + if (loadedTrigger) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + return; + } + // Group g = block % 16 = tid % 16: lanes l and l ^ 16 of every warp. + localMin = min(localMin, __shfl_xor_sync(0xffffffffu, localMin, 16)); + if (lane < 16) + atomicMin(&sm.groupMin[lane], localMin); + __syncthreads(); + // Every score of the row is loaded: release the sparse decode, so its launch and its pre-wait prologue overlap + // the rank below. This comes after this grid's own dependency wait, which the sparse decode relies on: it reads + // q and the seq_lens / block-table row before its own wait (sparseDecodeLc2.cu), safe only while every kernel + // between the producer and it releases its dependents after its own wait. + if (loadedTrigger) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + uint32_t bound = 0; +#pragma unroll + for (int g = 0; g < 16; ++g) + bound = max(bound, sm.groupMin[g]); +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + bool const isCand = block < valid && key[s] <= bound; + unsigned const ballot = __ballot_sync(0xffffffffu, isCand); + int base = 0; + if (lane == 0 && ballot) + base = atomicAdd(&sm.numCand, __popc(ballot)); + base = __shfl_sync(0xffffffffu, base, 0); + if (isCand) + { + int const i = base + __popc(ballot & ((1u << lane) - 1)); + if (i < 1024) + sm.cand[i] = (static_cast(key[s]) << 32) | static_cast(block); + } + } + __syncthreads(); + int const n = sm.numCand; + if (n <= 1024) + { + for (int i = tid; i < n; i += 1024) + { + unsigned long long const mine = sm.cand[i]; + int rank = 0; + for (int j = 0; j < n; ++j) + rank += sm.cand[j] < mine; + if (rank < 16) + sm.picked[rank] = mine; + } + return; + } + // More than 1024 keys tied at the bound (degenerate rows): rank every candidate against the whole row. +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + if (block < valid && key[s] <= bound) + { + int rank = 0; + for (int j = 0; j < valid; ++j) + { + uint32_t const other = orderedKey(applyForcing(rowBase[j * blockStride], j, initBlocks, localStart)); + rank += other < key[s] || (other == key[s] && j < block); + } + if (rank < 16) + sm.picked[rank] = (static_cast(key[s]) << 32) | static_cast(block); + } + } +} + +__global__ void __launch_bounds__(1024) m3_select_long2_kernel(float const* __restrict__ scores, long long headStride, + long long blockStride, long long queryStride, int const* __restrict__ nValid, int* __restrict__ out, + long long outQueryStride, long long outHeadStride, int numKvHeads, int numBlocks, int initBlocks, int localBlocks, + bool earlyTrigger, bool loadedTrigger) +{ + __shared__ SelectSmem sm; + int const row = blockIdx.x, query = row / numKvHeads, head = row % numKvHeads; + int const tid = threadIdx.x, lane = tid & 31; + if (tid < 16) + { + sm.groupMin[tid] = 0xffffffffu; + sm.picked[tid] = ~0ull; + } + if (tid == 0) + sm.numCand = 0; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); + // The sparse decode that follows stages its table-independent inputs before its own grid wait; past this grid's + // wait those inputs are final, so it may launch now rather than at this grid's exit. + if (earlyTrigger) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + int const raw = nValid[query]; + int const valid = max(0, min(raw, numBlocks)); + long long const localStart = max(static_cast(raw) - localBlocks, 0LL); + float const* rowBase = scores + head * headStride + query * queryStride; + __syncthreads(); + if (valid <= 2048) + selectRow<2>(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + else if (valid <= 4096) + selectRow<4>(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + else + selectRow(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + __syncthreads(); + if (tid >= 32) + return; + uint32_t blockKey = ~0u; + if (lane < 16) + { + unsigned long long const c = sm.picked[lane]; + if (c != ~0ull && static_cast(c >> 32) != kNegInfKey) + blockKey = static_cast(c & 0xffffffffu); + } +#pragma unroll + for (int size = 2; size <= 32; size <<= 1) + { +#pragma unroll + for (int stride = size >> 1; stride > 0; stride >>= 1) + { + uint32_t const other = __shfl_xor_sync(0xffffffffu, blockKey, stride); + bool const ascending = (lane & size) == 0, lower = (lane & stride) == 0; + blockKey = (lower == ascending) ? min(blockKey, other) : max(blockKey, other); + } + } + if (lane < 16) + out[query * outQueryStride + head * outHeadStride + lane] = blockKey == ~0u ? -1 : static_cast(blockKey); +} + +} // namespace + +cudaError_t invokeSelectLong2(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* out, int64_t outQueryStride, int64_t outHeadStride, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQ, int32_t initBlocks, int32_t localBlocks, bool usePdl, bool earlyTrigger, + int32_t triggerPhase, cudaStream_t stream) +{ + if (numBlocks > kSelectLong2MaxBlocks || triggerPhase < 0 || triggerPhase > 1) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(totalQ * numKvHeads); + cfg.blockDim = dim3(1024); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = usePdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_select_long2_kernel, scores, static_cast(headStride), + static_cast(blockStride), static_cast(queryStride), nValidBlocks, out, + static_cast(outQueryStride), static_cast(outHeadStride), numKvHeads, numBlocks, + initBlocks, localBlocks, earlyTrigger, !earlyTrigger && triggerPhase == 1); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h new file mode 100644 index 000000000000..9fba25f03c77 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +//! Largest row (blocks per query) the long-row top-16 select covers: 9 keys per thread of one 1024-thread CTA. +constexpr int32_t kSelectLong2MaxBlocks = 9216; + +//! MiniMax-M3 top-16 block select over per-KV-head scores [numKvHeads, numBlocks, totalQ] (element strides given), +//! one 1024-thread CTA per (query, KV head). Same contract as minimax_m3_select_blocks: init / local forcing, the +//! per-query valid-block count, ties to the smaller block, forced -inf picks emitted as -1, ascending ids with -1 +//! tail padding. Writes out[query * outQueryStride + head * outHeadStride + 0..15]. The SM107 TP2 copy of +//! minimaxM3/selectLong with a choice of where the dependent grid (the sparse decode) is released: at this grid's exit +//! (triggerPhase 0), right after every score of the row is loaded (1), or, with earlyTrigger, right after this grid's +//! dependency wait. Every release point follows this grid's own dependency wait. +cudaError_t invokeSelectLong2(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* out, int64_t outQueryStride, int64_t outHeadStride, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQ, int32_t initBlocks, int32_t localBlocks, bool usePdl, bool earlyTrigger, + int32_t triggerPhase, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu new file mode 100644 index 000000000000..cb7d83b89bb4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu @@ -0,0 +1,89 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr int kThreads = 128; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves). The +// instruction sequence of the Triton kernel's PTX at its defaults (host/silu/silu_and_mul_plain.ptx): g * -log2(e) as +// one product, ex2.approx, + 1, div.full 1 / that, g * s, (u + 0.0) as an IEEE add, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t siluMulBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 gl, gh, ul, uh, t0, t1, e0, e1, s0, s1, p0, p1, q0, q1, r0, r1, one, zero;\n\t" + ".reg .b16 hl, hh;\n\t" + "mov.b32 {hl, hh}, %1;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mul.rn.f32 t0, gl, 0fBFB8AA3B;\n\tmul.rn.f32 t1, gh, 0fBFB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 p0, gl, s0;\n\tmul.rn.f32 p1, gh, s1;\n\t" + "mov.b32 zero, 0f00000000;\n\tadd.rn.f32 q0, ul, zero;\n\tadd.rn.f32 q1, uh, zero;\n\t" + "mul.rn.f32 r0, q0, p0;\n\tmul.rn.f32 r1, q1, p1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +// Eight outputs per thread: 16-byte loads of the gate and the up half, one 16-byte store. +__global__ void __launch_bounds__(kThreads) + siluAndMulKernel(uint4 const* __restrict__ x, uint4* __restrict__ out, int64_t rows, int64_t vecs) +{ + int64_t const row = blockIdx.y; + int64_t const v = static_cast(blockIdx.x) * kThreads + threadIdx.x; + if (row >= rows || v >= vecs) + { + return; + } + uint4 const g = x[row * 2 * vecs + v]; + uint4 const u = x[row * 2 * vecs + vecs + v]; + uint4 o; + o.x = siluMulBf16x2(g.x, u.x); + o.y = siluMulBf16x2(g.y, u.y); + o.z = siluMulBf16x2(g.z, u.z); + o.w = siluMulBf16x2(g.w, u.w); + out[row * vecs + v] = o; +} + +} // namespace + +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream) +{ + int64_t const vecs = d / 8; + dim3 const grid(static_cast((vecs + kThreads - 1) / kThreads), static_cast(rows)); + siluAndMulKernel<<>>( + reinterpret_cast(x), reinterpret_cast(out), rows, vecs); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h new file mode 100644 index 000000000000..01b32137ac4a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h @@ -0,0 +1,40 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +//! out[m, j] = bf16(silu(g) * (u + 0)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both +//! contiguous, d a multiple of 8, 16-byte aligned rows): bit for bit the Triton silu_and_mul kernel at its defaults +//! (no limit, alpha 1, beta 0), the plain SwiGLU of the Eagle3 drafter MLP. +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu new file mode 100644 index 000000000000..d96f51d5d150 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu @@ -0,0 +1,640 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// MiniMax-M3 long-context sparse decode attention in one launch, SM107 TP2 geometry (32 query heads on 2 KV heads +// per rank; staircase copy of minimaxM3/sparseDecodeLc.cu with a KV-head grid dimension). +// +// Grid: (tokens * 16, kv heads), one cluster of 16 CTAs per (query token, KV head); CTA `r` attends the token's `r`-th +// selected block (the Triton producer's chunk `r` at 16 chunks), keeps the partial in shared memory, and after a +// cluster barrier merges head `r` from all 16 partials over DSMEM, writing the BF16 row and, optionally, its MXFP8 copy +// and swizzled UE8M0 scales exactly as the Triton merge's `_store_mxfp8_head` lays them out. +// +// Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +// PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +// log-domain recurrence. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeads = 16; +constexpr int kDim = 128; +constexpr int kPage = 128; +constexpr int kTopk = 16; +constexpr int kThreads = 128; +constexpr int kPitch = kDim + 8; // halves per smem row: 272 B, conflict-free ldmatrix +constexpr float kNegInf = -__builtin_huge_valf(); +constexpr int kBtCap = 4096; // block-table entries staged in smem by PREWAIT (512K tokens) + +using Params = MinimaxM3Sm107Tp2SparseDecodeLcParams; + +struct Smem +{ + __half qs[kHeads * kPitch]; + __half vs[kPage * kPitch]; + __half ps[kHeads * kPitch]; + float recv[kTopk][kDim]; // head `rank` partials pushed by the 16 CTAs, chunk-major + float recv_lse[kTopk]; + float wgt[kTopk]; // the 16 merge weights, computed once per CTA + float red[4][kHeads]; + float rowm[kHeads]; + float rowl[kHeads]; + alignas(8) unsigned long long rbar; // the split-K partials of head `rank` landed (complete_tx) +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t e2m1x2_to_f16x2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 s;\n cvt.u8.u32 s, %1;\n cvt.rn.f16x2.e2m1x2 %0, s;\n}" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t e4m3x2_to_f16x2(uint16_t pair) +{ + uint32_t out; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(pair)); + return out; +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ void ldsm_x4(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x4_t(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma(float (&c)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint32_t map_cluster(void const* local, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smem_u32(local)), "r"(rank)); + return remote; +} + +// Posted DSMEM stores that complete_tx on the receiving CTA's mbarrier `bar`. +__device__ __forceinline__ void st_async_v2(uint32_t addr, float a, float b, uint32_t bar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(addr), + "f"(a), "f"(b), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ void st_async_f32(uint32_t addr, float a, uint32_t bar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(addr), + "r"(__float_as_uint(a)), "r"(bar) + : "memory"); +} + +// One packed NVFP4 row (64 bytes, 128 values) with its 8 E4M3 scales -> 128 fp16 values in smem. +__device__ __forceinline__ void dequant_row(__half* dst, uint8_t const* packed, uint8_t const (&sc)[8], bool valid) +{ + uint4 const* src = reinterpret_cast(packed); + uint32_t* d = reinterpret_cast(dst); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + uint32_t words[4] = {w.x, w.y, w.z, w.w}; + // 16 bytes = 32 values = scale columns 2c and 2c + 1. +#pragma unroll + for (int h = 0; h < 2; ++h) + { + uint8_t const s = valid ? sc[2 * c + h] : 0; + uint32_t const sh = e4m3x2_to_f16x2(static_cast(s | (s << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); +#pragma unroll + for (int j = 0; j < 2; ++j) + { + uint32_t const word = words[2 * h + j]; +#pragma unroll + for (int b = 0; b < 4; ++b) + { + uint32_t const v2 = e2m1x2_to_f16x2((word >> (8 * b)) & 0xFF); + __half2 const x = __hmul2(*reinterpret_cast<__half2 const*>(&v2), scale); + d[c * 16 + h * 8 + j * 4 + b] = *reinterpret_cast(&x); + } + } + } + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +template +constexpr size_t smemBytes() +{ + return PREWAIT ? sizeof(Smem) + kBtCap * sizeof(int) : sizeof(Smem); +} + +template +__global__ void __launch_bounds__(kThreads) m3_lc_attn_kernel(__grid_constant__ Params const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(128) uint8_t smem_raw[]; + Smem& s = *reinterpret_cast(smem_raw); + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const tok = blockIdx.x / kTopk; + int const kvh = blockIdx.y; + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + + if (tid == 0) + { + // One arrival (here) plus every cluster CTA's partial row (kDim floats) and LSE for head `rank`. + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smem_u32(&s.rbar)) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(&s.rbar)), + "r"(kTopk * (kDim + 1) * 4) + : "memory"); + } + // Every peer must be running before its shared memory takes a DSMEM store: arrive now, wait before pushing. + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + + // Pre-wait read contract: everything that does not depend on the top-16 table is read before the grid dependency. + // - seq_lens: written inside the step's graph by msa_kv_update (the KV-length correction). That kernel is a plain + // launch, not a programmatic dependent, so it has completed before any later kernel of the step starts. + // - the query rows: written by the producer, two grids back. This holds only while every kernel between the + // producer and this grid (the long-row select) triggers its dependents after its own griddepcontrol.wait, so + // this grid cannot start before the producer completed. A kernel inserted into that chain, or a select that + // triggers before its wait, breaks it: these reads then move after the griddepcontrol.wait below. + int const req = tok / p.dql; + int const qoff = tok - req * p.dql; + int const seq = p.seq_lens[req]; + int const kv_len = max(seq - p.dql + qoff + 1, 0); + int const nblocks = (kv_len + kPage - 1) / kPage; + int const real_topk = min(kTopk, nblocks); + // q: 16 heads x 128 e4m3 -> fp16; thread = (head, 16-value chunk). + { + int const h = tid >> 3; + int const c = (tid & 7) * 16; + uint4 const w = *reinterpret_cast(p.q + tok * p.q_tok + (kvh * kHeads + h) * p.q_head + c); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; + uint32_t* d = reinterpret_cast(s.qs + h * kPitch + c); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + d[2 * i] = e4m3x2_to_f16x2(static_cast(words[i] & 0xFFFF)); + d[2 * i + 1] = e4m3x2_to_f16x2(static_cast(words[i] >> 16)); + } + } + + // PREWAIT: the request's block-table row (written before the graph runs) is staged before the grid dependency. + int* const s_bt = reinterpret_cast(smem_raw + sizeof(Smem)); + bool bt_smem = false; + if constexpr (PREWAIT) + { + int const nbt = (max(seq, 0) + kPage - 1) / kPage; + bt_smem = nbt <= kBtCap; + if (bt_smem) + { + for (int i = tid; i < nbt; i += kThreads) + s_bt[i] = __ldg(p.block_table + req * p.bt_stride + i); + } + } + // N1: the block this (KV head, row, slot) selected at the layer's previous decode step recurs ~2/3 of the time; + // pull its K/V page and scales into L2 while the select runs. The ids were written by this kernel one step back + // (complete before the graph's next replay) and the block-table row before the graph runs. Read and written by + // the same thread, so the record below cannot overtake this read. Out-of-range ids (stale rows) are skipped. + bool const n1 = p.n1_ids != nullptr && tok < p.n1_rows; + int* const n1_slot = n1 ? p.n1_ids + kvh * p.n1_h + tok * kTopk + rank : nullptr; + if (n1 && tid == 64) + { + int const nb = *n1_slot; + if (nb >= 0 && nb < nblocks) + { + int64_t const pg = __ldg(p.block_table + req * p.bt_stride + nb); + if (pg >= 0 && pg < p.n1_pages) + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.k + pg * p.k_blk + kvh * p.k_h), + "r"(kPage * 64) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.v + pg * p.v_blk + kvh * p.v_h), + "r"(kPage * 64) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.ksc + pg * p.ksc_blk + kvh * p.ksc_h), + "r"(kPage * 8) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.vsc + pg * p.vsc_blk + kvh * p.vsc_h), + "r"(kPage * 8) + : "memory"); + } + } + } + asm volatile("griddepcontrol.wait;" ::: "memory"); + if constexpr (EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } + int blk = static_cast(rank) < real_topk ? p.topk[kvh * p.t_h + tok * p.t_tok + rank * p.t_k] : -1; + if (n1 && tid == 64) + { + *n1_slot = blk; + } + bool const active = blk >= 0; + if constexpr (PREWAIT) + { + __syncthreads(); // s_bt complete + } + int64_t const page_pw = !active ? 0 : (bt_smem ? s_bt[blk] : p.block_table[req * p.bt_stride + blk]); + + float m_row[2] = {kNegInf, kNegInf}; + float acc[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + acc[j][i] = 0.F; + float l_row[2] = {0.F, 0.F}; + + // K stays packed in registers: the MMA B fragment of token n = lane / 4 at k-step s is the f16x2 pair of + // packed bytes 8s + q and 8s + 4 + q (q = lane % 4), both under scale column s. + uint32_t kw[4][16]; + uint2 ksc[4]; + int64_t page = 0; + if (active) + { + page = page_pw; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const n = warp * 32 + j * 8 + (lane >> 2); + bool const valid = blk * kPage + n < kv_len; + uint4 const* src = reinterpret_cast(p.k + page * p.k_blk + kvh * p.k_h + n * 64); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 const w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + kw[j][4 * c] = w.x; + kw[j][4 * c + 1] = w.y; + kw[j][4 * c + 2] = w.z; + kw[j][4 * c + 3] = w.w; + } + ksc[j] = valid ? __ldg(reinterpret_cast(p.ksc + page * p.ksc_blk + kvh * p.ksc_h + n * 8)) + : make_uint2(0, 0); + } + bool const valid = blk * kPage + tid < kv_len; + uint8_t vsc[8]; + if (valid) + { + uint8_t const* vb = p.vsc + page * p.vsc_blk + kvh * p.vsc_h + (tid >> 2) * 32 + (tid & 3); +#pragma unroll + for (int c = 0; c < 8; ++c) + vsc[c] = __ldg(vb + 4 * c); + } + dequant_row(s.vs + tid * kPitch, p.v + page * p.v_blk + kvh * p.v_h + tid * 64, vsc, valid); + } + __syncthreads(); + + if (active) + { + float const scale_log2 = p.sm_scale * 1.4426950409F * __ldg(p.kgs); + int const q4 = (lane & 3) * 8; + // S = q K^T for this warp's 32 tokens: rows g / g + 8, columns 2 * (lane % 4) + {0, 1} of each n-tile. + float sacc[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + sacc[j][i] = 0.F; +#pragma unroll + for (int ks = 0; ks < kDim / 16; ++ks) + { + uint32_t a[4]; + ldsm_x4(a, smem_u32(s.qs + (lane & 15) * kPitch + ks * 16 + (lane >> 4) * 8)); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + uint32_t const scw = ks < 4 ? ksc[j].x : ksc[j].y; + uint8_t const sc = static_cast(scw >> (8 * (ks & 3))); + uint32_t const sh = e4m3x2_to_f16x2(static_cast(sc | (sc << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); + uint32_t const r0 = e2m1x2_to_f16x2((kw[j][2 * ks] >> q4) & 0xFF); + uint32_t const r1 = e2m1x2_to_f16x2((kw[j][2 * ks + 1] >> q4) & 0xFF); + __half2 const h0 = __hmul2(*reinterpret_cast<__half2 const*>(&r0), scale); + __half2 const h1 = __hmul2(*reinterpret_cast<__half2 const*>(&r1), scale); + mma(sacc[j], a, *reinterpret_cast(&h0), *reinterpret_cast(&h1)); + } + } + int const g = lane >> 2; + float mloc[2] = {kNegInf, kNegInf}; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const t = warp * 32 + j * 8 + 2 * (lane & 3) + (i & 1); + float const v = blk * kPage + t < kv_len ? 0.F + sacc[j][i] * scale_log2 : kNegInf; + sacc[j][i] = v; + mloc[i >> 1] = fmaxf(mloc[i >> 1], v); + } +#pragma unroll + for (int r = 0; r < 2; ++r) + { + mloc[r] = fmaxf(mloc[r], __shfl_xor_sync(0xffffffff, mloc[r], 1)); + mloc[r] = fmaxf(mloc[r], __shfl_xor_sync(0xffffffff, mloc[r], 2)); + } + if ((lane & 3) == 0) + { + s.red[warp][g] = mloc[0]; + s.red[warp][g + 8] = mloc[1]; + } + __syncthreads(); + if (tid < kHeads) + s.rowm[tid] = fmaxf(fmaxf(s.red[0][tid], s.red[1][tid]), fmaxf(s.red[2][tid], s.red[3][tid])); + __syncthreads(); + m_row[0] = s.rowm[g]; + m_row[1] = s.rowm[g + 8]; + float lloc[2] = {0.F, 0.F}; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; i += 2) + { + float const p0 = ex2(sacc[j][i] - m_row[i >> 1]); + float const p1 = ex2(sacc[j][i + 1] - m_row[i >> 1]); + lloc[i >> 1] += p0 + p1; + int const row = g + (i >> 1) * 8; + int const t = warp * 32 + j * 8 + 2 * (lane & 3); + *reinterpret_cast<__half2*>(s.ps + row * kPitch + t) = __floats2half2_rn(p0, p1); + } +#pragma unroll + for (int r = 0; r < 2; ++r) + { + lloc[r] += __shfl_xor_sync(0xffffffff, lloc[r], 1); + lloc[r] += __shfl_xor_sync(0xffffffff, lloc[r], 2); + } + __syncthreads(); // rowm reads done before red is reused; ps complete + if ((lane & 3) == 0) + { + s.red[warp][g] = lloc[0]; + s.red[warp][g + 8] = lloc[1]; + } + // O[:, warp's 32 dims] = P V over all 128 tokens. +#pragma unroll + for (int kk = 0; kk < kPage; kk += 16) + { + uint32_t a[4]; + ldsm_x4(a, smem_u32(s.ps + (lane & 15) * kPitch + kk + (lane >> 4) * 8)); +#pragma unroll + for (int jp = 0; jp < 2; ++jp) + { + uint32_t b[4]; + int const t = kk + (lane & 7) + ((lane >> 3) & 1) * 8; + int const d = warp * 32 + jp * 16 + (lane >> 4) * 8; + ldsm_x4_t(b, smem_u32(s.vs + t * kPitch + d)); + mma(acc[2 * jp], a, b[0], b[1]); + mma(acc[2 * jp + 1], a, b[2], b[3]); + } + } + __syncthreads(); + if (tid < kHeads) + s.rowl[tid] = (s.red[0][tid] + s.red[1][tid]) + (s.red[2][tid] + s.red[3][tid]); + __syncthreads(); + l_row[0] = s.rowl[g]; + l_row[1] = s.rowl[g + 8]; + } + + // Normalize the partial and push head h's row with its LSE (log2) into CTA h's recv[rank] slot, as the Triton + // producer stores chunk `rank`; st.async stores that complete_tx on CTA h's barrier. + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + { + int const g = lane >> 2; + float const vgs = __ldg(p.vgs); +#pragma unroll + for (int r = 0; r < 2; ++r) + { + float scale, lse; + if constexpr (LINEAR) + { + scale = l_row[r] > 0.F ? 1.0F / l_row[r] : 0.F; + lse = l_row[r] > 0.F ? m_row[r] + lg2(l_row[r]) : kNegInf; + } + else + { + float const lse_i = active ? m_row[r] + lg2(ex2(kNegInf - m_row[r]) + l_row[r]) : kNegInf; + scale = lse_i > kNegInf ? ex2(m_row[r] - lse_i) : 0.F; + lse = lse_i; + } + scale = scale * vgs; + int const row = g + r * 8; + uint32_t const dst = map_cluster(&s.recv[rank][0], row); + uint32_t const rb = map_cluster(&s.rbar, row); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const d = warp * 32 + j * 8 + 2 * (lane & 3); + st_async_v2(dst + d * 4, acc[j][2 * r] * scale, acc[j][2 * r + 1] * scale, rb); + } + if (warp == 0 && (lane & 3) == 0) + st_async_f32(map_cluster(&s.recv_lse[rank], row), lse, rb); + } + } + { + // Every CTA's partial row and LSE for head `rank` landed (complete_tx on this CTA's barrier). + uint32_t done = 0; + while (!done) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(smem_u32(&s.rbar)) + : "memory"); + } + } + + // CTA `rank` merges head `rank` across the 16 partials; thread = output dim. The weights are the Triton merge's + // (exp2 of each chunk's LSE minus the largest, divided by their sum in chunk order), formed once per CTA: thread + // c < 16 takes chunk c's exponential, then divides it by the sum every one of those threads forms in the same + // order. + int const h = static_cast(rank); + if (tid < kTopk) + { + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + lse_max = fmaxf(lse_max, s.recv_lse[c]); + lse_max = lse_max == kNegInf ? 0.F : lse_max; + s.red[0][tid] = ex2(s.recv_lse[tid] - lse_max); + } + __syncthreads(); + if (tid < kTopk) + { + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + denom += s.red[0][c]; + denom = denom > 0.F ? denom : 1.F; + s.wgt[tid] = __fdiv_rn(s.red[0][tid], denom); + } + __syncthreads(); + float o = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + o += s.recv[c][tid] * s.wgt[c]; + __nv_bfloat16 const ob = __float2bfloat16_rn(o); + p.out[tok * p.out_stride + (kvh * kHeads + h) * kDim + tid] = ob; + + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + float const xq = x * qs; + uint16_t q2; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2) : "f"(0.F), "f"(xq)); + p.q8[static_cast(tok) * (p.num_sf_k_tiles * 128) + (kvh * kHeads + h) * kDim + tid] + = static_cast(q2 & 0xFF); + if (lane == 0) + { + int const k_sf = (kvh * kHeads + h) * 4 + warp; + int64_t const off = static_cast(tok / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (tok % 32) * 16 + ((tok % 128) / 32) * 4 + k_sf % 4; + p.sf[off] = static_cast(e); + } + } + if constexpr (!EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +cudaError_t configure() +{ + auto kernel = m3_lc_attn_kernel; + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes())); + if (err != cudaSuccess) + { + return err; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); +} + +template +cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) +{ + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs + // every variant eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + auto kernel = m3_lc_attn_kernel; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(num_tokens * kTopk, prm.num_kv_heads); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smemBytes(); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kTopk; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +} // namespace + +cudaError_t invokeMinimaxM3Sm107Tp2SparseDecodeLc(MinimaxM3Sm107Tp2SparseDecodeLcParams const& params, + int32_t numTokens, bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream) +{ + if (prewait) + { + if (linearSoftmax) + { + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); + } + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); + } + if (linearSoftmax) + { + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); + } + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h new file mode 100644 index 000000000000..b20a1894f255 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h @@ -0,0 +1,77 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Operands of the MiniMax-M3 one-launch long-context sparse decode, SM107 TP2 geometry. Strides are in elements of +//! the pointed-to type. Geometry is fixed: 16 query heads per KV head, head_dim 128, page 128, top-k 16. +struct MinimaxM3Sm107Tp2SparseDecodeLcParams +{ + uint8_t const* q; // e4m3 [T, kv_heads * 16, 128] + int64_t q_tok, q_head; + uint8_t const* k; // packed e2m1 [pages, kv_heads, 128, 64] + uint8_t const* v; + int64_t k_blk, v_blk, k_h, v_h; + uint8_t const* ksc; // e4m3 [pages, kv_heads, 128, 8], token-major + uint8_t const* vsc; // e4m3, 4x4 token-quad order + int64_t ksc_blk, vsc_blk, ksc_h, vsc_h; + float const* kgs; + float const* vgs; + int const* topk; // [kv_heads, T, 16] + int64_t t_h, t_tok, t_k; + int const* block_table; + int64_t bt_stride; + int const* seq_lens; + int dql; + float sm_scale; + __nv_bfloat16* out; // [T, kv_heads * 2048] + int64_t out_stride; + uint8_t* q8; // e4m3 [rows, num_sf_k_tiles * 128] + uint8_t* sf; + int num_sf_k_tiles; + int num_kv_heads; + // N1 (optional): this layer's selected block id per (KV head, query row, slot) at its previous decode step, + // int32 [kv_heads, n1_rows, 16]. Before the grid dependency each CTA pulls into L2 the K/V page and scales of the + // block its slot selected last step, then records this step's block. A hint: no value changes. + int* n1_ids; + int64_t n1_h; + int n1_rows; + int64_t n1_pages; +}; + +//! One cluster of 16 CTAs per (query token, KV head): CTA r attends the token's r-th selected block of that head, +//! then the cluster merges the 16 partials over DSMEM into the BF16 rows (and, when mxfp8, the MXFP8 copy with +//! swizzled UE8M0 scales). prewait stages the request's block-table row in shared memory before the grid dependency +//! wait (rows longer than 4096 blocks read it from global memory); the arithmetic is the same in both variants. +//! The device code is SM100-family only; on other architectures the kernel compiles to an empty body. +cudaError_t invokeMinimaxM3Sm107Tp2SparseDecodeLc(MinimaxM3Sm107Tp2SparseDecodeLcParams const& params, + int32_t numTokens, bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu new file mode 100644 index 000000000000..22c073f74b3c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr int kThreads = 256; + +__global__ void __launch_bounds__(kThreads) + specOffsetsKernel(int32_t* __restrict__ pos, int32_t const* __restrict__ posOff, int64_t nPos, + int32_t* __restrict__ kv, int32_t const* __restrict__ kvOff, int64_t nKv) +{ + for (int64_t i = threadIdx.x; i < nPos; i += kThreads) + { + pos[i] = static_cast(static_cast(pos[i]) + static_cast(posOff[i])); + } + for (int64_t j = threadIdx.x; j < nKv; j += kThreads) + { + kv[j] = static_cast(static_cast(kv[j]) + static_cast(kvOff[j])); + } +} + +} // namespace + +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream) +{ + specOffsetsKernel<<<1, kThreads, 0, stream>>>(pos, posOff, nPos, kv, kvOff, nKv); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h new file mode 100644 index 000000000000..1de212a37179 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +//! The overlap scheduler's speculative step correction in one launch: pos[i] += posOff[i] for i < nPos and +//! kv[j] += kvOff[j] for j < nKv (int32, wrapping as the ATen adds do). A plain launch: it is a step-input writer that +//! pre-wait readers rely on (rubin-20261015/PREWAIT_AUDIT.md item b). +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu new file mode 100644 index 000000000000..9b6dbb42faed --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr int kThreads = 512; +constexpr int kWarps = kThreads / 32; + +// One warp per copy. Plain (coherent) loads: a device copy may read bytes another warp of this CTA just stored, which +// the CTA barrier between the two kinds makes visible; a non-coherent load could return the stale line. +__device__ __forceinline__ void copyWarp(uint8_t* dst, uint8_t const* src, int64_t bytes, int lane) +{ + auto const d = reinterpret_cast(dst); + auto const s = reinterpret_cast(src); + if (((d | s | static_cast(bytes)) & 15) == 0) + { + auto* d16 = reinterpret_cast(dst); + auto const* s16 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 16; i += 32) + { + d16[i] = s16[i]; + } + } + else if (((d | s | static_cast(bytes)) & 3) == 0) + { + auto* d4 = reinterpret_cast(dst); + auto const* s4 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 4; i += 32) + { + d4[i] = s4[i]; + } + } + else + { + for (int64_t i = lane; i < bytes; i += 32) + { + dst[i] = src[i]; + } + } +} + +// One CTA: a step's window is a few dozen copies of a few hundred bytes, so the launch is the cost; a CTA barrier +// orders the host copies before the device copies where the parallel scatter needed a second launch. +__global__ void __launch_bounds__(kThreads) stepPrologueKernel(uint8_t* window, int64_t* epoch, int64_t maxDesc) +{ + auto* header = reinterpret_cast(window); + int64_t const numHost = min(header[0], maxDesc); + int64_t const numDevice = min(header[1], maxDesc); + int64_t const* desc = header + 2; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (epoch != nullptr && threadIdx.x == 0) + { + *epoch += 1; + } + for (int64_t i = warp; i < numHost; i += kWarps) + { + int64_t const* c = desc + 3 * i; + copyWarp(reinterpret_cast(c[0]), window + c[1], c[2], lane); + } + __syncthreads(); + for (int64_t i = warp; i < numDevice; i += kWarps) + { + // The device copies' slots follow all header[0] host copies, as the parallel scatter indexes them. + int64_t const* c = desc + 3 * (header[0] + i); + copyWarp(reinterpret_cast(c[0]), reinterpret_cast(c[1]), c[2], lane); + } + // Every thread has read the header and the descriptors: consume the window. + __syncthreads(); + if (threadIdx.x < 2) + { + header[threadIdx.x] = 0; + } +} + +} // namespace + +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream) +{ + stepPrologueKernel<<<1, kThreads, 0, stream>>>(window, epoch, maxDesc); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h new file mode 100644 index 000000000000..44a7b7bcb44d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h @@ -0,0 +1,47 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +//! The decode graph's step-input prologue in one launch (step_inputs.StepInputPack with early_copy): the work of the +//! two parallel scatter launches and the header clear that follow them, in that order. +//! +//! The window is the pack's staging layout (kernels/stepInputsScatter.h): an int64 header `[numHostCopies, +//! numDeviceCopies]`, then one `(dst, src, bytes)` int64 triple per copy, host copies first; a host copy's `src` is a +//! byte offset into the window, a device copy's a device address. Up to `maxDesc` copies of each kind are written (the +//! slots the parallel scatter's grid covers), every host copy before any device copy, so a device copy may read a +//! destination a host copy filled; the copies of one kind must not overlap (the pack issues the window only then). +//! Afterwards the header is zeroed, so a replay with no new upload writes nothing. With `epoch` non-null, it is +//! incremented once. +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu new file mode 100644 index 000000000000..5da59fcec01a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr int kThreads = 256; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves): the +// instruction sequence of the Triton kernel at limit 7, alpha 1.702, beta 1 (Indexer kdev17, bitwise on SM107): the +// gate clamp as min.bf16x2 before the widening, the up clamp as min.xorsign.abs, (u + 1), (g * -1.702) * log2(e), +// ex2.approx, + 1, div.full 1 / that, s * g, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t swigluOaiBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, %1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +template +__global__ void __launch_bounds__(kThreads) + swigluOaiKernel(uint32_t const* __restrict__ x, uint32_t* __restrict__ out, int64_t rows, int64_t halfD) +{ + if constexpr (kPdl) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int64_t const i = blockIdx.x * static_cast(kThreads) + threadIdx.x; + if (i >= rows * halfD) + { + return; + } + int64_t const row = i / halfD; + int64_t const col = i - row * halfD; + uint32_t const* xr = x + row * 2 * halfD; + out[i] = swigluOaiBf16x2(xr[col], xr[halfD + col]); +} + +} // namespace + +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream) +{ + int64_t const n = rows * halfD; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast((n + kThreads - 1) / kThreads)); + cfg.blockDim = dim3(kThreads); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return pdl ? cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD) + : cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h new file mode 100644 index 000000000000..c8db72a5efd3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +//! The dense MLP's swigluoai activation: out[m, j] = bf16(min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + +//! 1)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both contiguous, d even): bit for bit +//! the Triton silu_and_mul kernel at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With `pdl` the launch allows +//! programmatic serialization and the grid waits on its dependency, then releases its dependents before it computes +//! (the down GEMV stages its weights during the activation and reads it only after its own grid wait). +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu new file mode 100644 index 000000000000..e173c91016c1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu @@ -0,0 +1,296 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// The TP4 dense layers' QK norm + RoPE with FP8 output. The norm and RoPE are the generic fusedQKNormRopeKernel's at +// head_dim 128 without YaRN or mRoPE, expression for expression; the store rounds Q/K to BF16 and then to E4M3 +// (saturating), the two roundings of the stock BF16 kernel followed by its FP8 cast. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace common +{ +// tensorrt_llm/common/envUtils.h (not included: its CUDA utility headers do not build under the extension's +// no-half-conversion flags). +bool getEnvEnablePDL(); +} // namespace common + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace denseQkNormRopeFp8 +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kVecSize = kElemsPerThread * sizeof(__nv_bfloat16) / 4; +constexpr int kBlockSize = 256; +static_assert(kVecSize == 2, "a lane loads 4 BF16 values as one uint2"); + +// tensorrt_llm::common::warpReduceSum: the same butterfly, so the same sum. +__device__ __forceinline__ float warpReduceSum(float val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + val = val + __shfl_xor_sync(0xffffffff, val, mask, 32); + } + return val; +} + +__device__ __forceinline__ void storeFp8ViaBf16(__nv_fp8_e4m3* out, long long offsetThread, float const (&elements)[4]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + float2 const rounded = __bfloat1622float2(__float22bfloat162_rn(make_float2(elements[i], elements[i + 1]))); + __nv_fp8x2_e4m3 const packed(rounded); + reinterpret_cast<__nv_fp8x2_storage_t*>(&out[offsetThread])[i / 2] = packed.__x; + } +} + +// The K (kv 0) or V (kv 1) row of one head of token tokenIdx in the HND paged pool, at the token's cache slot. +__device__ __forceinline__ __nv_fp8_e4m3* cacheRow(KvCacheView const& cache, int kv, int headIdx, int tokenIdx) +{ + long long const slot = static_cast(cache.outCacheLoc[tokenIdx]); + long long const page = slot / cache.tokensPerBlock; + long long const within = slot % cache.tokensPerBlock; + return cache.base + page * cache.pageStride + kv * cache.kvStride + headIdx * cache.headStride + + within * cache.tokenStride; +} + +template +__global__ void denseQkNormRopeFp8Kernel(__nv_bfloat16 const* qkvIn, __nv_fp8_e4m3* qkvOut, int numHeadsQ, + int numHeadsK, int numHeadsV, int rotaryDim, float eps, __nv_bfloat16 const* qWeight, __nv_bfloat16 const* kWeight, + float base, int const* positionIds, int numTokens, bool useGemma, KvCacheView cache) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + // PDL: every read of this step's q/k/v rows, positions and cache tables follows the wait. The dependent (the + // full-context decode) reads this kernel's outputs only after its own griddepcontrol.wait, so it is released at + // entry and stages its CTAs while this kernel runs. + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId; + int const totalQkHeads = numHeadsQ + numHeadsK; + int const totalHeads = totalQkHeads + numHeadsV; + int const tokenIdx = globalWarpIdx / totalHeads; + int const localHeadIdx = globalWarpIdx % totalHeads; + if (tokenIdx >= numTokens) + { + return; + } + bool const isQ = localHeadIdx < numHeadsQ; + bool const isV = localHeadIdx >= totalQkHeads; + int headIdx; + int segStart; + if (isQ) + { + headIdx = localHeadIdx; + segStart = 0; + } + else if (!isV) + { + headIdx = localHeadIdx - numHeadsQ; + segStart = numHeadsQ * kHeadDim; + } + else + { + headIdx = localHeadIdx - totalQkHeads; + segStart = totalQkHeads * kHeadDim; + } + + using vec_T = uint2; + int const offsetThread + = tokenIdx * totalHeads * kHeadDim + segStart + headIdx * kHeadDim + laneId * kElemsPerThread; + float elements[kElemsPerThread]; + float sumOfSquares = 0.0f; + { + vec_T vec = *reinterpret_cast(&qkvIn[offsetThread]); +#pragma unroll + for (int i = 0; i < kVecSize; i++) + { + float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&vec) + i)); + sumOfSquares += vals.x * vals.x; + sumOfSquares += vals.y * vals.y; + elements[2 * i] = vals.x; + elements[2 * i + 1] = vals.y; + } + } + if (isV) + { + // V is BF16 already: its single E4M3 rounding is the stock cast's. + storeFp8ViaBf16(qkvOut, offsetThread, elements); + if (kWriteCache && tokenIdx < cache.numLiveTokens && cache.outCacheLoc[tokenIdx] >= 0) + { + storeFp8ViaBf16(cacheRow(cache, 1, headIdx, tokenIdx), laneId * kElemsPerThread, elements); + } + return; + } + + sumOfSquares = warpReduceSum(sumOfSquares); + float rms_rcp = rsqrtf(sumOfSquares / static_cast(kHeadDim) + eps); + for (int i = 0; i < kElemsPerThread; i++) + { + int dim = laneId * kElemsPerThread + i; + float weight = isQ ? __bfloat162float(qWeight[dim]) : __bfloat162float(kWeight[dim]); + elements[i] *= rms_rcp * (useGemma ? (1.0f + weight) : weight); + } + + float elements2[kElemsPerThread]; + float cos_vals[kElemsPerThread]; + float sin_vals[kElemsPerThread]; + float const neg2_log2base_over_rd = -2.0f * __log2f(base) / static_cast(rotaryDim); + int const rd_mask = rotaryDim - 1; + bool const rd_is_pow2 = ((rotaryDim & rd_mask) == 0); + float const pos_id = static_cast(positionIds[tokenIdx]); + if constexpr (kInterleave) + { + for (int i = 0; i < kElemsPerThread; i++) + { + if (i % 2 == 0) + { + elements2[i] = -elements[i + 1]; + } + else + { + elements2[i] = elements[i - 1]; + } + int dim_idx = laneId * kElemsPerThread + i; + int half_dim = dim_idx / 2; + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); + } + } + else + { + __syncwarp(); + int pairOffset = (rotaryDim / 2) / kElemsPerThread; + for (int i = 0; i < kElemsPerThread; i++) + { + elements2[i] = __shfl_xor_sync(0xffffffff, elements[i], pairOffset); + if (laneId < pairOffset) + { + elements2[i] = -elements2[i]; + } + int dim_idx = laneId * kElemsPerThread + i; + dim_idx = rd_is_pow2 ? ((dim_idx * 2) & rd_mask) : ((dim_idx * 2) % rotaryDim); + int half_dim = dim_idx / 2; + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); + } + __syncwarp(); + } + + // attention_factor is 1 without YaRN: (x * cos + x2 * sin) * 1 rounds as the generic kernel's product does. + float const attention_factor = 1.0f; + if (rotaryDim == kHeadDim) + { + for (int i = 0; i < kElemsPerThread; i++) + { + elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor; + } + } + else + { + for (int i = 0; i < kElemsPerThread; i++) + { + int dim_idx = laneId * kElemsPerThread + i; + if (dim_idx < rotaryDim) + { + elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor; + } + } + } + storeFp8ViaBf16(qkvOut, offsetThread, elements); + if constexpr (kWriteCache) + { + // Rows past the live count (padded graph rows) own no slot, nor does a negative (sentinel) slot. + if (!isQ && tokenIdx < cache.numLiveTokens && cache.outCacheLoc[tokenIdx] >= 0) + { + storeFp8ViaBf16(cacheRow(cache, 0, headIdx, tokenIdx), laneId * kElemsPerThread, elements); + } + // The compact query [tokens, q heads, 128] the decode reads, beside the packed row. + if (isQ && cache.qCompact != nullptr) + { + storeFp8ViaBf16(cache.qCompact, + (static_cast(tokenIdx) * numHeadsQ + headIdx) * kHeadDim + laneId * kElemsPerThread, + elements); + } + } +} + +} // namespace + +cudaError_t launch(void const* qkvIn, void* qkvOut, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, + int rotaryDim, float eps, void const* qWeight, void const* kWeight, float base, bool interleave, + int const* positionIds, bool useGemma, KvCacheView const* cache, cudaStream_t stream) +{ + int const totalWarps = numTokens * (numHeadsQ + numHeadsK + numHeadsV); + int const warpsPerBlock = kBlockSize / 32; + dim3 const grid((totalWarps + warpsPerBlock - 1) / warpsPerBlock); + auto const* in = static_cast<__nv_bfloat16 const*>(qkvIn); + auto* out = static_cast<__nv_fp8_e4m3*>(qkvOut); + auto const* qw = static_cast<__nv_bfloat16 const*>(qWeight); + auto const* kw = static_cast<__nv_bfloat16 const*>(kWeight); + KvCacheView const view = cache != nullptr ? *cache : KvCacheView{}; + auto const run = [&](auto kernel) -> cudaError_t + { + // The maximum shared-memory carveout lets the dependent's large-smem CTAs co-reside on the SMs this small + // kernel occupies once it releases them. + cudaError_t const err = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (err != cudaSuccess) + { + return err; + } + cudaLaunchConfig_t config = {}; + config.gridDim = grid; + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + return cudaLaunchKernelEx(&config, kernel, in, out, numHeadsQ, numHeadsK, numHeadsV, rotaryDim, eps, qw, kw, + base, positionIds, numTokens, useGemma, view); + }; + if (cache != nullptr) + { + return interleave ? run(denseQkNormRopeFp8Kernel) : run(denseQkNormRopeFp8Kernel); + } + return interleave ? run(denseQkNormRopeFp8Kernel) : run(denseQkNormRopeFp8Kernel); +} + +} // namespace denseQkNormRopeFp8 +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h new file mode 100644 index 000000000000..1049033b07c9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace denseQkNormRopeFp8 +{ + +constexpr int kHeadDim = 128; + +//! A layer's FP8 K/V pool in the HND paged layout [pages, 2 (K, V), kv heads, tokens per block, 128] (element +//! strides), with the step's per-token cache slots (page * tokensPerBlock + within) for its first numLiveTokens rows. +struct KvCacheView +{ + __nv_fp8_e4m3* base = nullptr; + int const* outCacheLoc = nullptr; + long long pageStride = 0; + long long kvStride = 0; + long long headStride = 0; + long long tokenStride = 0; + int tokensPerBlock = 1; + int numLiveTokens = 0; + //! Optional: every row's Q also stored compact, [tokens, q heads, 128]. + __nv_fp8_e4m3* qCompact = nullptr; +}; + +//! The MiniMax-M3 TP4 dense layers' per-head QK RMSNorm + partial RoPE of a packed BF16 [tokens, (q + k + v) * 128] +//! row, written out of place as FP8 E4M3 with V copy-cast. Q and K are rounded to BF16 before E4M3, as the stock path +//! (the BF16 fused_qk_norm_rope output cast to FP8) rounds them. With cache, the K and V rows of the live tokens are +//! also stored into the paged pool at their slots (the bytes of the qkvOut rows). +//! +//! PDL (with TRTLLM_ENABLE_PDL): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction, so +//! nothing is read or written before the producer has completed, and the kernel releases its dependent right after +//! it. The dependent may read only step-invariant data before its own griddepcontrol.wait. +cudaError_t launch(void const* qkvIn, void* qkvOut, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, + int rotaryDim, float eps, void const* qWeight, void const* kWeight, float base, bool interleave, + int const* positionIds, bool useGemma, KvCacheView const* cache, cudaStream_t stream); + +} // namespace denseQkNormRopeFp8 +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu new file mode 100644 index 000000000000..2e67b3ae7ac1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu @@ -0,0 +1,667 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h" + +#include +#include +#include + +// One pass of the fused greedy Eagle3 draft loop of MiniMax-M3 (GB300, TP4, <= 8 requests per pass). +// +// One kernel replaces the draft lm_head GEMM, the logits all-gather / float cast / argmax of the in-tree greedy draft +// sampling and the draft loop's per-pass bookkeeping, at the in-tree precision: +// +// * The vocab-sharded BF16 lm_head weight is stored in mma.m16n8k16 A-fragment order (the same values): per 16-row +// tile and k16 step, 512 bytes; lane (g, j) holds 16 bytes = {row g k 2j..2j+1, row g+8 k 2j.., row g k 2j+8.., +// row g+8 k 2j+8..}. +// * Or, by default, in a lossless 12-bit coding of those fragments (75% of the bytes; the target head's coding, see +// targetHead.cu): per 16-row tile and unit of 4 k16 steps a lane reads three 16-byte words, the sign/mantissa bytes +// of steps 0-1, of steps 2-3, then 4 code words (exponent - base, 4 bits; 15 = an exception). The decoded values +// are the checkpoint's BF16 bits; an exception decodes to +0 and the segment's flush adds the exact product x * w +// from the tile's list (products of BF16 values are exact in FP32). +// * A persistent grid (one 16-warp CTA per SM) streams whole tiles over the full K with BF16 MMAs and FP32 +// accumulation; the BF16 head input rows sit in shared memory in B-fragment order, unchanged. +// * Each logit is rounded to BF16 (the in-tree logits are the BF16 GEMM output) and becomes a 64-bit key +// ordered(value) << 32 | ~global_index, with -0 taken as +0: the larger value wins and ties pick the lowest index, as +// torch.max / torch.argmax over the gathered logits do. Keys are max-reduced per row: warp shuffles, then shared +// memory, then atomicMax. +// * The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping +// (Lamport: 0 = empty, a key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the +// slots, and runs the pass's bookkeeping. Call n uses slot half n % 2. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace draftStep +{ +namespace +{ +constexpr int kUnroll = 8; // k16 chunks in flight per warp +constexpr int kChunkBytes = 512; // one tile's k16 step: 32 lanes x 16 bytes + +// Coded units (4 k16 steps, 3 x 16 bytes per lane) in flight per warp: 3 up to 4 rows, 2 at 8 rows. +template +constexpr int codedUnroll() +{ + return MP >= 8 ? 2 : 3; +} + +__device__ __forceinline__ void mma_bf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// BF16 pair p (values 2p, 2p + 1) of a lane's 8 coded values: sign/mantissa bytes sm (4 per word), codes word cw. +__device__ __forceinline__ uint32_t coded_decode_pair(uint32_t sm_word, int byte0, uint32_t cw, int p, uint32_t base) +{ + uint32_t const b = __byte_perm(sm_word, 0u, byte0 == 0 ? 0x4140 : 0x4342); // {b0, 0, b1, 0} + uint32_t const q = (cw >> (4 * p)) & 0x000f000fu; + uint32_t const e = q * 128u + base; // exponent field of both halves + // Code 15 (an exception, sign/mantissa byte 0) decodes to +0. + return ((b & 0x007f007fu) | ((b << 8) & 0x80008000u) | e) & __vcmpne2(q, 0x000f000fu); +} + +// One coded unit (4 k16 steps) against the BF16 B fragments: steps alternate between the two accumulators, as the +// k16 steps of the plain BF16 head do. +template +__device__ __forceinline__ void chunk_mma_coded(uint4 const& sm01, uint4 const& sm23, uint4 const& cw, uint32_t base, + uint2 const* xs, int unit, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint32_t const smw[8] = {sm01.x, sm01.y, sm01.z, sm01.w, sm23.x, sm23.y, sm23.z, sm23.w}; + uint32_t const cws[4] = {cw.x, cw.y, cw.z, cw.w}; +#pragma unroll + for (int ci = 0; ci < 4; ++ci) + { + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[((unit * 4 + ci) * 4 + j) * MP + g]; + uint32_t const w0 = coded_decode_pair(smw[2 * ci], 0, cws[ci], 0, base); + uint32_t const w1 = coded_decode_pair(smw[2 * ci], 2, cws[ci], 1, base); + uint32_t const w2 = coded_decode_pair(smw[2 * ci + 1], 0, cws[ci], 2, base); + uint32_t const w3 = coded_decode_pair(smw[2 * ci + 1], 2, cws[ci], 3, base); + if (ci & 1) + mma_bf16(db, w0, w1, w2, w3, b.x, b.y); + else + mma_bf16(da, w0, w1, w2, w3, b.x, b.y); + } +} + +// Exceptions of `tile` with k in [kb, ke): v += x * w (exact) for the owner lanes (row g or g + 8; tokens 2j, +// 2j + 1 < M), x from the staged B fragments. Warp-uniform (shuffles). +template +__device__ __forceinline__ void coded_exceptions( + CodedArgs const& cd, int tile, int kb, int ke, uint2 const* xs, int M, int lane, float (&v)[4]) +{ + int const g = lane >> 2, j = lane & 3; + int const e0 = __ldg(cd.exc_off + tile), e1 = __ldg(cd.exc_off + tile + 1); + for (int b0 = e0; b0 < e1; b0 += 32) + { + uint32_t pos = 0xffffffffu; + float2 wv = make_float2(0.f, 0.f); + if (b0 + lane < e1) + { + pos = __ldg(cd.exc_pos + b0 + lane); + wv = __ldg(cd.exc_w + b0 + lane); + } + int const cnt = min(32, e1 - b0); + for (int q = 0; q < cnt; ++q) + { + uint32_t const pq = __shfl_sync(0xffffffffu, pos, q); + float const w0 = __shfl_sync(0xffffffffu, wv.x, q); + int const r = static_cast(pq >> 16), k = static_cast(pq & 0xffffu); + if (k >= kb && k < ke && (r & 7) == g) + { + int const st = k >> 4, rr = k & 15; + bool const hi = r >= 8; +#pragma unroll + for (int q2 = 0; q2 < 2; ++q2) + { + int const n = 2 * j + q2; + if (n < M) + { + uint2 const bw = xs[(st * 4 + ((rr & 7) >> 1)) * MP + n]; + uint32_t const word = rr >= 8 ? bw.y : bw.x; + float const xv = __uint_as_float((rr & 1) ? (word & 0xffff0000u) : (word << 16)); + float const d = __fmul_rn(xv, w0); + if (hi) + v[2 + q2] += d; + else + v[q2] += d; + } + } + } + __syncwarp(); + } + } +} + +__device__ __forceinline__ uint4 ldg_stream(uint4 const* p) +{ + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.u32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +// The argmax value of a logit: its BF16 rounding (the in-tree logits' dtype), -0 as +0. +__device__ __forceinline__ float key_value(float v) +{ + float const r = __bfloat162float(__float2bfloat16(v)); + return r == 0.f ? 0.f : r; +} + +__device__ __forceinline__ unsigned long long make_key(float v, uint32_t index) +{ + uint32_t u = __float_as_uint(v); + u = (u & 0x80000000u) ? ~u : (u | 0x80000000u); + return (static_cast(u) << 32) | static_cast(0xFFFFFFFFu - index); +} + +__device__ __forceinline__ unsigned long long shfl_max(unsigned long long k, int o) +{ + unsigned long long const ok = __shfl_xor_sync(0xffffffffu, k, o); + return ok > k ? ok : k; +} + +// One k16 step of a tile: the lane's A fragment w against the activation B fragment of step `chunk` (rows g < MP). +// Even and odd steps accumulate separately (two independent MMA chains); the tile's sum is da + db. +template +__device__ __forceinline__ void chunk_mma( + uint4 const& w, uint2 const* xs, int chunk, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[(chunk * 4 + j) * MP + g]; + if (chunk & 1) + mma_bf16(db, w.x, w.y, w.z, w.w, b.x, b.y); + else + mma_bf16(da, w.x, w.y, w.z, w.w, b.x, b.y); +} + +// A warp's pieces of the tiles it shares with other warps (its first and last local tile), kept for the warp-order sum. +struct Pieces +{ + float4 head = make_float4(0.f, 0.f, 0.f, 0.f); + float4 tail = make_float4(0.f, 0.f, 0.f, 0.f); +}; + +// Whether the warp's chunk range [q_begin, q_end) covers the CTA's whole range of local tile tl (the CTA streams +// chunks [lo, lo + span) of its local tiles). +__device__ __forceinline__ bool covers_tile(int tl, int kchunks, int lo, int span, int q_begin, int q_end) +{ + return q_begin <= max(tl * kchunks, lo) && q_end >= min((tl + 1) * kchunks, lo + span); +} + +// Local tile t's sum (D-FIX): the accumulator when one warp streamed the CTA's whole range of it, else its pieces +// from the warps that streamed it, added in warp order from 0 (each warp's first tile is its head piece, any later +// one its tail piece). Deterministic: the order is fixed by the launch shape, and with two pieces it is the same +// sum as the atomics' (two addends from 0 commute). +__device__ __forceinline__ float4 ordered_tile( + float4 const* acc_t, float4 const* stage, int t, int kchunks, int lo, int span, int lane) +{ + int const r_lo = max(t * kchunks, lo); + int const r_last = min((t + 1) * kchunks, lo + span) - 1; + int w_lo = 0, w_hi = 0; + for (int w = 1; w < kWarps; ++w) + { + int const qb = lo + w * span / kWarps; + if (qb <= r_lo) + w_lo = w; + if (qb <= r_last) + w_hi = w; + } + if (w_lo == w_hi) + return acc_t[lane]; + float4 s = make_float4(0.f, 0.f, 0.f, 0.f); + for (int w = w_lo; w <= w_hi; ++w) + { + int const qb = lo + w * span / kWarps; + float4 const x = stage[(w * 2 + (qb / kchunks == t ? 0 : 1)) * 32 + lane]; + s.x += x.x; + s.y += x.y; + s.z += x.z; + s.w += x.w; + } + return s; +} + +// Add a warp's segment of a tile into the CTA's accumulator (with a coded head, the segment's exceptions first); +// with `fixed` (D-FIX) a shared tile's piece goes to `pc` instead (ordered_tile adds it). +template +__device__ __forceinline__ void flush_tile(float4* acc, float (&da)[4], float (&db)[4], Params const& p, int tile, + int kb, int ke, uint2 const* xs, int lane, Pieces& pc, bool fixed, bool whole, bool first) +{ + float v[4] = {da[0] + db[0], da[1] + db[1], da[2] + db[2], da[3] + db[3]}; + if constexpr (CODED) + coded_exceptions(p.cd, tile, kb, ke, xs, p.M, lane, v); + if (fixed) + { + // A piece covering the CTA's whole range of the tile is its only addend (0 + v, as the atomic gave); a + // warp's other pieces (at most its first and last tile) wait in registers for the warp-order sum. + if (whole) + *acc = make_float4(0.f + v[0], 0.f + v[1], 0.f + v[2], 0.f + v[3]); + else if (first) + pc.head = make_float4(v[0], v[1], v[2], v[3]); + else + pc.tail = make_float4(v[0], v[1], v[2], v[3]); + } + else + { + float* a = reinterpret_cast(acc); + atomicAdd(a + 0, v[0]); + atomicAdd(a + 1, v[1]); + atomicAdd(a + 2, v[2]); + atomicAdd(a + 3, v[3]); + } +#pragma unroll + for (int e = 0; e < 4; ++e) + { + da[e] = 0.f; + db[e] = 0.f; + } +} + +// Balanced streaming GEMV over the fragment-ordered weight, persistent grid of 16-warp CTAs. CTA c owns tiles +// [t_begin, t_end); its (tile, chunk) space is split into 16 equal contiguous ranges, one per warp, so every warp +// streams the same number of bytes. A tile straddles at most two warps; their partial sums are added into shared +// memory (two addends: the order does not change the sum). MP = rows padded to 1 / 2 / 4 / 8. A chunk is one k16 +// step of the BF16 head, or one coded unit (4 k16 steps) of the coded head (CODED: no shared-memory prologue). +template +__global__ void __launch_bounds__(kWarps * 32, 1) + draft_step_kernel(Params const p, int kchunks, int pf_chunks, int pf_off) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(16) uint2 xs[]; // [K/16 steps][4 (j)][MP (g)], then the weight prologue + __shared__ float4 acc_s[kMaxTilesPerCta][32]; + __shared__ unsigned long long red[kWarps][kMaxTokens]; + __shared__ __align__(8) unsigned long long pf_bar[kWarps]; + __shared__ bool s_last; + __shared__ int s_tok[kMaxTokens]; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const g = lane >> 2; + int const j = lane & 3; + constexpr int kChunkK = CODED ? 64 : 16; // k per chunk + constexpr int kWords = CODED ? 3 : 1; // 16-byte words per lane and chunk + constexpr int KU = CODED ? codedUnroll() : kUnroll; + int const K = kchunks * kChunkK; + int const M = p.M; + int const t_begin = static_cast((static_cast(blockIdx.x) * p.n_tiles) / gridDim.x); + int const t_end = static_cast((static_cast(blockIdx.x + 1) * p.n_tiles) / gridDim.x); + int const n_local = t_end - t_begin; + int const span = n_local * kchunks; + int const q_begin = warp * span / kWarps; + int const q_end = (warp + 1) * span / kWarps; + long long const base = static_cast(t_begin) * kchunks; // first chunk of this CTA + uint4 const* const wf = p.wf; + + // Weights are static: start this warp's stream before waiting on the producer of x. The warp's first pf_chunks + // chunks come in by one bulk copy into shared memory (dynamic offset pf_off), which fills the grid-dependency + // wait and the staging; the register ring starts after them. + int const n_pf = CODED ? 0 : min(pf_chunks, q_end - q_begin); + uint4* const wpf = reinterpret_cast(reinterpret_cast(xs) + pf_off) + warp * pf_chunks * 32; + if (n_pf > 0 && lane == 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + uint32_t const bytes = static_cast(n_pf) * kChunkBytes; + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + static_cast(__cvta_generic_to_shared(wpf))), + "l"(wf + (base + q_begin) * 32), "r"(bytes), "r"(bar) + : "memory"); + } + __syncwarp(); + int const q_reg = q_begin + n_pf; + uint4 buf[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q_reg + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = ldg_stream(wf + ((base + q_reg + u) * kWords + l) * 32 + lane); + } + } + for (int i = threadIdx.x; i < n_local * 32; i += blockDim.x) + acc_s[i / 32][i % 32] = make_float4(0.f, 0.f, 0.f, 0.f); + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // Stage the head input rows in B-fragment order: item (k16 step S, row gg) = two 16-byte loads of x, four + // fragment words {x[2jj], x[2jj + 1]}, {x[2jj + 8], x[2jj + 9]} (BF16 pairs, moved unchanged). Rows >= M are 0. + constexpr int kSU = MP >= 4 ? 4 : 2; + int const n_items = (K / 16) * MP; + for (int i0 = threadIdx.x; i0 < n_items; i0 += kSU * blockDim.x) + { + uint4 xv[kSU][2]; +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + int const gg = i % MP; + int const S = i / MP; + xv[u][0] = make_uint4(0u, 0u, 0u, 0u); + xv[u][1] = make_uint4(0u, 0u, 0u, 0u); + if (i < n_items && gg < M) + { + long long const xr = p.x_rows != nullptr ? p.x_rows[gg] : gg; + uint4 const* xp = reinterpret_cast(p.x + xr * K + 16 * S); + xv[u][0] = xp[0]; + xv[u][1] = xp[1]; + } + } +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + if (i >= n_items) + continue; + int const gg = i % MP; + int const S = i / MP; + uint32_t const lo[4] = {xv[u][0].x, xv[u][0].y, xv[u][0].z, xv[u][0].w}; + uint32_t const hi[4] = {xv[u][1].x, xv[u][1].y, xv[u][1].z, xv[u][1].w}; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) + xs[(S * 4 + jj) * MP + gg] = make_uint2(lo[jj], hi[jj]); + } + } + __syncthreads(); + + bool const b_active = g < MP; + float da[4] = {0.f, 0.f, 0.f, 0.f}; + float db[4] = {0.f, 0.f, 0.f, 0.f}; + int tl = q_begin / kchunks; // local tile of the current accumulator + int c = q_begin - tl * kchunks; // chunk within that tile + int cs = c; // first chunk of the warp's segment of that tile + int const tl0 = tl; // the warp's first local tile (its head piece) + Pieces pc; + if (n_pf > 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + asm volatile( + "{\n .reg .pred p;\n PF_WAIT%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n" + " @!p bra PF_WAIT%=;\n}\n" ::"r"(bar) + : "memory"); + for (int q = q_begin; q < q_reg; ++q) + { + chunk_mma(wpf[(q - q_begin) * 32 + lane], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, 0, span, q_begin, q_end), tl == tl0); + c = 0; + cs = 0; + ++tl; + } + } + } + for (int q = q_reg; q < q_end; q += KU) + { + uint4 nxt[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + KU + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + nxt[u][l] = ldg_stream(wf + ((base + q + KU + u) * kWords + l) * 32 + lane); + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + u < q_end) + { + if constexpr (CODED) + chunk_mma_coded(buf[u][0], buf[u][kWords > 1 ? 1 : 0], buf[u][kWords > 2 ? 2 : 0], p.cd.base, + xs, c, j, g, b_active, da, db); + else + chunk_mma(buf[u][0], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, 0, span, q_begin, q_end), tl == tl0); + c = 0; + cs = 0; + ++tl; + } + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = nxt[u][l]; + } + if (c != 0) + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, c * kChunkK, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, 0, span, q_begin, q_end), tl == tl0); + __syncthreads(); + // D-FIX: stage every warp's head / tail piece. xs and the weight prologue are dead after the barrier above (the + // launcher sizes them to hold the 16 x 2 x 32 pieces). + float4* const stage = reinterpret_cast(xs); + if (p.fixedOrder) + { + stage[(warp * 2) * 32 + lane] = pc.head; + stage[(warp * 2 + 1) * 32 + lane] = pc.tail; + __syncthreads(); + } + + // Accumulator lane (g, j) of a tile: rows g / g + 8 (x, y / z, w) of tokens 2j / 2j + 1 (x, z / y, w). + unsigned long long best0 = 0ull, best1 = 0ull; // tokens 2j, 2j + 1 + for (int t = warp; t < n_local; t += kWarps) + { + float4 const v = p.fixedOrder ? ordered_tile(acc_s[t], stage, t, kchunks, 0, span, lane) : acc_s[t][lane]; + int const r0 = (t_begin + t) * 16 + g; + uint32_t const i0 = static_cast(p.vocab_offset + r0); + unsigned long long k0 = make_key(key_value(v.x), i0); + unsigned long long const k0b = make_key(key_value(v.z), i0 + 8); + unsigned long long k1 = make_key(key_value(v.y), i0); + unsigned long long const k1b = make_key(key_value(v.w), i0 + 8); + k0 = k0b > k0 ? k0b : k0; + k1 = k1b > k1 ? k1b : k1; + best0 = k0 > best0 ? k0 : best0; + best1 = k1 > best1 ? k1 : best1; + } +#pragma unroll + for (int o = 4; o < 32; o <<= 1) + { + best0 = shfl_max(best0, o); + best1 = shfl_max(best1, o); + } + asm volatile("griddepcontrol.launch_dependents;"); + if (lane < 4) + { + red[warp][2 * lane] = best0; + red[warp][2 * lane + 1] = best1; + } + __syncthreads(); + if (threadIdx.x < M) + { + unsigned long long b = 0ull; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + b = red[w][threadIdx.x] > b ? red[w][threadIdx.x] : b; + if (b != 0ull) + atomicMax(p.keys + threadIdx.x, b); + } + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned int const prev = atomicAdd(p.counters, 1u); + s_last = prev == gridDim.x - 1; + } + __syncthreads(); + if (!s_last) + return; + + __threadfence(); + unsigned int const call = *reinterpret_cast(p.counters + 1); + int const half = static_cast(call & 1u); + if (threadIdx.x < M) + { + int const t = threadIdx.x; + unsigned long long const mine = atomicExch(p.keys + t, 0ull); + unsigned long long* slot_mc = p.ws_mc + (half * kMaxTokens + t) * p.world; + unsigned long long* slot_uc = p.ws_uc + (half * kMaxTokens + t) * p.world; + slot_mc[p.rank] = mine; + unsigned long long best = 0ull; + for (int r = 0; r < p.world; ++r) + { + unsigned long long v; + do + { + v = *reinterpret_cast(slot_uc + r); + } while (v == 0ull); + best = v > best ? v : best; + slot_uc[r] = 0ull; + } + s_tok[t] = static_cast(0xFFFFFFFFu - static_cast(best & 0xFFFFFFFFull)); + if (p.out_keys != nullptr) + p.out_keys[t] = best; + } + __syncthreads(); + if (threadIdx.x == 0) + { + p.counters[0] = 0u; + p.counters[1] = call + 1u; + } + if (threadIdx.x < M) + { + StepArgs const& st = p.st; + int const b = threadIdx.x; + int const tok = s_tok[b]; + st.draft_out[b * st.draft_stride + st.step] = tok; + st.ids_next[b] = tok; + long long const r = p.x_rows != nullptr ? p.x_rows[b] : b; + st.pos_next[b] = st.pos_src[r] + 1; + if (st.step == 0 && st.num_acc != nullptr) + { + st.kv_lens[b] -= st.draft_len - st.num_acc[b]; + st.seq_lens[b] = 1; + } + else if (st.step > 0) + { + st.kv_lens[b] += 1; + } + } + if (p.hasFold) + { + // The fold tail (H-R2-3) after this pass's tokens and bookkeeping: the barrier makes this CTA's writes above + // visible to warp 0; every earlier operand was written by grids this one waited on. All CTAs triggered their + // dependents before the exchange, which is safe: a dependent's griddepcontrol.wait waits for this whole grid + // (this tail included) to complete, so only a dependent that read the tail's outputs before its own wait + // could race, and none does (the next step's readers run in the next graph replay). + __syncthreads(); + if (threadIdx.x < 32) + { + foldTailWarp(p.fold, static_cast(threadIdx.x)); + } + } +#endif +} + +template +cudaError_t launch(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + auto kern = draft_step_kernel; + size_t const xs_bytes = static_cast(K / 16) * 4 * MP * sizeof(uint2); + // Shared-memory weight prologue: chunks per warp that fit next to the kernel's static + dynamic shared memory. + static int const sStaticSmem = [] + { + cudaFuncAttributes fa{}; + return cudaFuncGetAttributes(&fa, draft_step_kernel) == cudaSuccess + ? static_cast(fa.sharedSizeBytes) + : -1; + }(); + static int const sMaxSmemPerBlock = [] + { + int dev = 0, maxSmem = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + return maxSmem; + }(); + int pf = 0; + if (!CODED && sStaticSmem >= 0) + { + long long const room + = static_cast(sMaxSmemPerBlock) - sStaticSmem - static_cast(xs_bytes) - 1024; + long long const want = static_cast(kPrefetchKb) * 1024; + pf = static_cast(std::max(0ll, std::min(room, want)) / (kWarps * kChunkBytes)); + } + // At least the D-FIX piece staging (16 warps x head / tail x 32 lanes of float4), which reuses this space. + size_t const smem = std::max(xs_bytes + static_cast(pf) * kWarps * kChunkBytes, + static_cast(kWarps) * 2 * 32 * sizeof(float4)); + cudaError_t err = cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + return err; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(grid)); + cfg.blockDim = dim3(kWarps * 32); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, kern, p, K / (CODED ? 64 : 16), pf, static_cast(xs_bytes)); +} + +template +cudaError_t dispatchRows(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + if (p.M <= 1) + return launch<1, CODED>(p, K, grid, pdl, stream); + if (p.M <= 2) + return launch<2, CODED>(p, K, grid, pdl, stream); + if (p.M <= 4) + return launch<4, CODED>(p, K, grid, pdl, stream); + return launch<8, CODED>(p, K, grid, pdl, stream); +} + +} // namespace + +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + bool const coded = p.cd.exc_off != nullptr; + if (p.M < 1 || p.M > kMaxTokens || K % (coded ? 64 : 16) != 0 || grid < 1) + return cudaErrorInvalidValue; + return coded ? dispatchRows(p, K, grid, pdl, stream) : dispatchRows(p, K, grid, pdl, stream); +} + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h new file mode 100644 index 000000000000..85296a3fd6d7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h @@ -0,0 +1,107 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +//! One pass of the fused greedy Eagle3 draft loop of the MiniMax-M3 NVFP4 TP4 decode path (sm_103a): the BF16 +//! vocab-sharded draft head, the cross-rank argmax and the draft loop's per-pass bookkeeping in one kernel. See +//! draftStep.cu. +namespace draftStep +{ + +constexpr int kMaxTokens = 8; // head rows (requests) of one launch +constexpr int kWarps = 16; // warps per CTA +constexpr int kMaxTilesPerCta = 32; // 16-row weight tiles one CTA can own +//! Shared-memory weight prologue per CTA, in KB (capped by the free shared memory). +constexpr int kPrefetchKb = 112; + +// Draft-pass epilogue (last CTA, one thread per request row b): stores the token as draft token `step`, stages the +// next pass's input id and position, and advances the draft KV lengths the way the in-tree draft loop does between +// passes (pass 0 with num_acc: drop the rejected verify rows and switch to one token per request; later passes: +1). +struct StepArgs +{ + int* draft_out; // [nreq, draft_stride]: column `step` receives the token + int draft_stride; + int step; + int* ids_next; // [nreq] + int const* pos_src; // positions of the rows the head read (indexed like x) + int* pos_next; // [nreq] = pos_src[x_rows[b] or b] + 1 + int* kv_lens; // [nreq] draft KV lengths, updated in place + int const* num_acc; // pass 0: accepted counts (nullptr: no rollback) + int* seq_lens; // pass 0: set to 1 + int draft_len; +}; + +// The 12-bit coded head (see draftStep.cu): the per-tile exception lists and the exponent base. exc_off == nullptr: +// the head is the plain BF16 fragment-ordered copy. +struct CodedArgs +{ + int const* exc_off; // [n_tiles + 1] exception list offsets + uint32_t const* exc_pos; // row in tile << 16 | k, per exception + float2 const* exc_w; // {w, 0} per exception: the exact weight + uint32_t base; // exponent base << 7 in both BF16 halves (base << 7 | base << 23) +}; + +// Everything a draft-step launch takes (see invokeDraftStep). +struct Params +{ + __nv_bfloat16 const* x; // [*, K] head input rows + long long const* x_rows; // optional [M]: head input row of request b (pass 0: the last accepted verify row) + int M; // requests, <= kMaxTokens + uint4 const* wf; // [n_tiles][K / 16][32 lanes] 16-byte BF16 A fragments, or the 12-bit coded stream + // [n_tiles][K / 64][3][32 lanes] when cd.exc_off is set (see draftStep.cu) + int n_tiles; + int vocab_offset; // global id of this rank's first head row + unsigned long long* keys; + unsigned int* counters; + unsigned long long* ws_uc; // this rank's exchange slots [2][kMaxTokens][world] + unsigned long long* ws_mc; // their multicast address + int rank; + int world; + unsigned long long* out_keys; // optional [M]: the winning global key of each row + StepArgs st; + CodedArgs cd; + // The step's forward-tail store + next-step staging (foldTail.h), run by the last CTA after this pass's + // tokens and bookkeeping: set on the drafter's last pass only (dispatch.DRAFT_STEP_FOLD_TAIL). + FoldTailParams fold{}; + bool hasFold = false; + // D-FIX: sum the pieces of a tile that several warps stream in warp order (deterministic) instead of by + // shared-memory atomics in arrival order (dispatch.DRAFT_STEP_FIXED_ORDER). + bool fixedOrder = false; +}; + +//! Launch the draft step over `grid` 512-thread CTAs (K = the hidden size, a multiple of 16; of 64 for a coded head). +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream); + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu new file mode 100644 index 000000000000..a506548c1aec --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu @@ -0,0 +1,189 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// The TP4 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's GPT-NeoX +// path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin negated, +// dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale E4M3 +// store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query +// saturates the same way (448, not NaN). + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace drafterRopeFp8Kv +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kBlockSize = 256; + +template +__device__ __forceinline__ float rotate(float c, float x, float s, float pair) +{ + if constexpr (kRopeMode == 0) + { + return __fadd_rn(__fmul_rn(c, x), __fmul_rn(s, pair)); + } + else if constexpr (kRopeMode == 1) + { + return __fmaf_rn(c, x, __fmul_rn(s, pair)); + } + else + { + return __fmaf_rn(s, pair, __fmul_rn(c, x)); + } +} + +__device__ __forceinline__ void storeE4m3Pairs(__nv_fp8_e4m3* dst, __nv_bfloat16 const (&v)[kElemsPerThread]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + __nv_bfloat162 pair; + pair.x = v[i]; + pair.y = v[i + 1]; + __nv_fp8x2_e4m3 const packed(pair); + reinterpret_cast<__nv_fp8x2_storage_t*>(dst)[i / 2] = packed.__x; + } +} + +template +__global__ void __launch_bounds__(kBlockSize) drafterRopeFp8KvKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const lane = threadIdx.x % 32; + int const warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; + int const slots = p.numHeadsQ + p.numHeadsKv; + int const token = warp / slots; + if (token >= p.numTokens) + { + return; + } + int const slot = warp % slots; + bool const isQ = slot < p.numHeadsQ; + int const head = isQ ? slot : slot - p.numHeadsQ; + int const request = token / p.qLen; + int const row = token % p.qLen; + int const cacheLen = p.seqLens[request]; + int const kvPos = cacheLen - p.qLen + row; + int const ropePos + = p.specPositionOffsets != nullptr ? p.specPositionOffsets[request * p.qLen + row] + cacheLen - p.qLen : kvPos; + int const rowWidth = (p.numHeadsQ + 2 * p.numHeadsKv) * kHeadDim; + __nv_bfloat16 const* src + = p.qkv + static_cast(token) * rowWidth + (isQ ? head * kHeadDim : (p.numHeadsQ + head) * kHeadDim); + int const half = p.rotaryDim / 2; + float2 const* coef = p.cosSin + static_cast(ropePos) * half; + + __nv_bfloat16 rotated[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + int const d = lane * kElemsPerThread + i; + bool const firstHalf = d < half; + float const x = __bfloat162float(src[d]); + float const pair = __bfloat162float(src[firstHalf ? d + half : d - half]); + float2 const cs = d < p.rotaryDim ? coef[firstHalf ? d : d - half] : make_float2(1.0f, 0.0f); + float const s = firstHalf ? -cs.y : cs.y; + rotated[i] = __float2bfloat16(rotate(cs.x, x, s, pair)); + } + + if (isQ) + { + storeE4m3Pairs( + p.qOut + (static_cast(token) * p.numHeadsQ + head) * kHeadDim + lane * kElemsPerThread, rotated); + return; + } + + int const block = kvPos / p.tokensPerBlock; + long long const inPage = (static_cast(head) * p.tokensPerBlock + kvPos % p.tokensPerBlock) * kHeadDim + + lane * kElemsPerThread; + long long const pageElems = static_cast(p.numHeadsKv) * p.tokensPerBlock * kHeadDim; + int const* table = p.blockTable + request * p.blockTableRequestStride; + long long const kPage = table[block]; + long long const vPage = table[p.blockTableKvStride + block]; + storeE4m3Pairs(p.pool + kPage * pageElems + inPage, rotated); + + __nv_bfloat16 const* vSrc = p.qkv + static_cast(token) * rowWidth + + (p.numHeadsQ + p.numHeadsKv + head) * kHeadDim + lane * kElemsPerThread; + __nv_bfloat16 v[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + v[i] = vSrc[i]; + } + storeE4m3Pairs(p.pool + vPage * pageElems + inPage, v); +} + +//! With PDL and maxSharedCarveout, the maximum shared-memory carveout lets the full-context decode's large-smem CTAs +//! (the dependent this kernel releases right after its wait) co-reside on the SMs this small kernel occupies. Set once +//! per instantiation. +template +cudaError_t launchWithMaxSharedCarveout( + cudaLaunchConfig_t const& config, Params const& params, bool pdl, bool maxSharedCarveout) +{ + auto const kernel = drafterRopeFp8KvKernel; + if (pdl && maxSharedCarveout) + { + static cudaError_t const carveout = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (carveout != cudaSuccess) + { + return carveout; + } + } + return cudaLaunchKernelEx(&config, kernel, params); +} + +} // namespace + +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream) +{ + int const warps = params.numTokens * (params.numHeadsQ + params.numHeadsKv); + int const warpsPerBlock = kBlockSize / 32; + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((warps + warpsPerBlock - 1) / warpsPerBlock); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + switch (params.ropeMode) + { + case 0: return launchWithMaxSharedCarveout<0>(config, params, pdl, maxSharedCarveout); + case 1: return launchWithMaxSharedCarveout<1>(config, params, pdl, maxSharedCarveout); + case 2: return launchWithMaxSharedCarveout<2>(config, params, pdl, maxSharedCarveout); + default: return cudaErrorInvalidValue; + } +} + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h new file mode 100644 index 000000000000..2a39d03d9d76 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace drafterRopeFp8Kv +{ + +constexpr int kHeadDim = 128; + +struct Params +{ + //! [tokens, (numHeadsQ + 2 * numHeadsKv) * 128] BF16: the drafter layer's packed Q, K, V rows. + __nv_bfloat16 const* qkv; + //! [positions, rotaryDim / 2] (cos, sin) pairs: the layer's rotary_cos_sin table. + float2 const* cosSin; + //! [requests] cache lengths including this step's rows. + int const* seqLens; + //! [requests, qLen] spec-decoding position offsets of a multi-token pass, or null. + int const* specPositionOffsets; + //! [tokens, numHeadsQ, 128] E4M3 query out. + __nv_fp8_e4m3* qOut; + //! Base of the layer's K / V sub-pages, [*, numHeadsKv, tokensPerBlock, 128] E4M3. + __nv_fp8_e4m3* pool; + //! [requests, 2, maxBlocks] sub-page indices: K in row 0, V in row 1. + int const* blockTable; + long long blockTableRequestStride; + long long blockTableKvStride; + int numTokens; + int qLen; + int numHeadsQ; + int numHeadsKv; + int rotaryDim; + int tokensPerBlock; + //! How the rotation's two products are summed: 0 two products then the add, 1 fma(cos, x, sin * pair), + //! 2 fma(sin, pair, cos * x). + int ropeMode; +}; + +//! The MiniMax-M3 TP4 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, +//! rounded to BF16 (the in-tree preprocess's q output and cache input), then Q, K and V to saturating E4M3: K / V +//! as the in-tree preprocess stores them at unit KV scale, Q as torch's CUDA cast of the BF16 query. Rows map to +//! (request token / qLen, row token % qLen) at cache position seqLen - qLen + row. +//! +//! PDL (pdl): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction and the kernel +//! releases its dependent right after it. The dependent may read only step-invariant data before its own wait. +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream); + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu new file mode 100644 index 000000000000..e6a4f0da2039 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +namespace +{ + +// Row i of the step (lane c = column), as the Triton fold tail's program i: the store's padded rows, then the next +// step's inputs read from them (zero past each source width). +__global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) +{ + int const i = blockIdx.x; + if (i >= *p.ctrl) + { + return; + } + foldTailRow(p, i, static_cast(threadIdx.x)); +} + +__global__ void __launch_bounds__(32) foldTailOneWarpKernel(FoldTailParams const p) +{ + foldTailWarp(p, static_cast(threadIdx.x)); +} + +} // namespace + +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream) +{ + foldTailKernel<<>>(p); + return cudaGetLastError(); +} + +cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream) +{ + foldTailOneWarpKernel<<<1, 32, 0, stream>>>(p); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h new file mode 100644 index 000000000000..54a58b0569e0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h @@ -0,0 +1,129 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +//! Operands of the speculative step's forward-tail store and next-step staging (kernels/m3_fold_tail.py). Every tensor +//! is int32 except `slots` (int64); strides are in elements. The store views are [slots, width], the sources +//! [rows, width]. +struct FoldTailParams +{ + int32_t const* ctrl; // ctrl[0]: the rows the step armed (0: write nothing) + int64_t const* slots; + int32_t const* srcNew; + int64_t sn0, sn1; + int32_t snW; + int32_t* dstNew; + int64_t dn0, dn1; + int32_t dnW; + int32_t const* srcNext; + int64_t sx0, sx1; + int32_t sxW; + int32_t* dstNext; + int64_t dx0, dx1; + int32_t dxW; + int32_t const* srcDraft; + int64_t sd0, sd1; + int32_t sdW; + int32_t* dstDraft; + int64_t dd0, dd1; + int32_t ddW; + int32_t const* srcLens; + int32_t* dstLens; + int32_t* ids; // [rows * T] + int32_t* draftIds; // [rows * (T - 1)] + int32_t* posOff; // [rows * T] + int32_t* kvOff; // [rows] + int32_t tokensPerStep; +}; + +//! One CTA of 32 lanes per row, `rows` CTAs; the first ctrl[0] rows write. Widths and tokensPerStep must be <= 32. +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream); + +//! The same writes from ONE warp (one CTA of 32 lanes looping over the rows, foldTailWarp): the in-kernel form a caller +//! such as the drafter's last-pass draft_step runs in its last CTA; launched alone it is that form's oracle. +cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream); + +#ifdef __CUDACC__ +//! Row i of the tail at lane c (0..31), the work of foldTailKernel's CTA i, for a caller that runs the tail inside its +//! own kernel (e.g. the drafter's last-pass draft_step in its last CTA, after its final tokens are written and fenced). +//! The caller runs rows i < *p.ctrl (one warp per row, or one warp looping over the rows); same code as the kernel, so +//! the writes are bit for bit the standalone tail's. +__device__ __forceinline__ void foldTailRow(FoldTailParams const& p, int i, int c) +{ + int64_t const s = p.slots[i]; + int32_t const nw = c < p.snW ? p.srcNew[i * p.sn0 + c * p.sn1] : 0; + if (c < p.dnW) + { + p.dstNew[s * p.dn0 + c * p.dn1] = nw; + } + int32_t const nxt = c < p.sxW ? p.srcNext[i * p.sx0 + c * p.sx1] : 0; + if (c < p.dxW) + { + p.dstNext[s * p.dx0 + c * p.dx1] = nxt; + } + int32_t const drf = c < p.sdW ? p.srcDraft[i * p.sd0 + c * p.sd1] : 0; + if (c < p.ddW) + { + p.dstDraft[s * p.dd0 + c * p.dd1] = drf; + } + int32_t const ln = p.srcLens[i]; + int32_t const t = p.tokensPerStep; + int32_t const d = t - 1; + if (c == 0) + { + p.dstLens[s] = ln; + p.kvOff[i] = ln - t; + } + if (c < t) + { + p.ids[i * t + c] = c < p.dxW ? nxt : 0; + p.posOff[i * t + c] = ln; + } + if (c < d) + { + p.draftIds[i * d + c] = c < p.ddW ? drf : 0; + } +} + +//! The whole tail from one warp of a running kernel (lane = threadIdx.x % 32): rows i < *p.ctrl in order. +__device__ __forceinline__ void foldTailWarp(FoldTailParams const& p, int lane) +{ + int const rows = *p.ctrl; + for (int i = 0; i < rows; ++i) + { + foldTailRow(p, i, lane); + } +} +#endif + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu new file mode 100644 index 000000000000..2937d461854c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu @@ -0,0 +1,268 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// The split-KV merge for the MiniMax-M3 NVFP4 TP2 (Rubin) full-context decode. The producer is the PrimTS split-KV +// decode with its separate-reduction scratch: per (request, KV head, split, row) one normalized 16-bit partial O and +// one fp32 log2-LSE. Active splits per row follow the producer's split ranges (the causal union of the row's grouped +// Q CTA, instruction-group aligned): splits past them hold nothing and are not read. +// +// One CTA per (request, KV head, kRowsPerCta rows); per row kGroups split groups x kThreadsPerRow threads x 8 dims. +// A thread issues the loads of every split it owns (strided over the groups) before using any, folds them under its +// group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to +// BF16 once. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace fullCtxReduce +{ +namespace +{ + +constexpr int kGroups = 8; +constexpr int kThreadsPerRow = kHeadDim / 8; +constexpr int kRowsPerCta = 1; +constexpr int kSplitsPerGroup = kMaxSplits / kGroups; +constexpr int kThreads = kGroups * kThreadsPerRow * kRowsPerCta; + +__device__ int activeSplits(Params const& p, int b, int row) +{ + int const seqKv = p.seqLensKv[b]; + int visible = seqKv; + if (p.causal) + { + int const token = row / p.headsPerKv; + int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; + int const end = min(base + p.qTokensPerCta, p.seqLenQ); + visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); + } + int const tiles = (visible + p.tileKv - 1) / p.tileKv; + int const tilesPerGroup = p.splits * p.instsKv; + int const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; + int const local = max(groups * p.instsKv, p.instsKv); + return (tiles + local - 1) / local; +} + +template +__device__ float2 unpack(uint32_t word) +{ + if constexpr (kBf16Partial) + { + return __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&word)); + } + else + { + return __half22float2(*reinterpret_cast<__half2 const*>(&word)); + } +} + +template +__global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); + // The dependent (the o_proj activation quantization / GEMV) reads `out` only after its own griddepcontrol.wait, + // which waits for this whole grid, so it is released here and runs its weight-only prologue under the merge. + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const b = blockIdx.z; + int const h = blockIdx.y; + int const localRow = threadIdx.x / (kGroups * kThreadsPerRow); + int const row = blockIdx.x * kRowsPerCta + localRow; + int const g = (threadIdx.x / kThreadsPerRow) % kGroups; + int const t = threadIdx.x % kThreadsPerRow; + bool const validRow = row < p.rows; + int const active = validRow ? activeSplits(p, b, row) : 0; + + __shared__ float sMax[kRowsPerCta][kGroups]; + __shared__ float sSum[kRowsPerCta][kGroups]; + __shared__ float sAcc[kRowsPerCta][kGroups][kHeadDim]; + + uint4 const* partialO = static_cast(p.partialO); + size_t const kvIdx = static_cast(b) * p.numKvHeads + h; + float lse[kSplitsPerGroup]; + uint4 vec[kSplitsPerGroup]; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + int const s = g + i * kGroups; + if (s < active) + { + size_t const r = (kvIdx * p.splits + s) * p.rows + row; + lse[i] = __ldcg(p.partialStats + r); + vec[i] = __ldcg(partialO + r * kThreadsPerRow + t); + } + } + float groupMax = -CUDART_INF_F; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active) + { + groupMax = fmaxf(groupMax, lse[i]); + } + } + // A split of the CTA's causal union can hold no key of this row (log2-LSE -inf). Such a split, and a group of + // only such splits, is skipped rather than weighted 0: its partial O carries no defined value. + float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + float sum = 0.f; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active && lse[i] != -CUDART_INF_F) + { + float const w = exp2f(lse[i] - groupMax); + sum += w; + uint32_t const* words = reinterpret_cast(&vec[i]); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = unpack(words[j]); + acc[2 * j] += w * f.x; + acc[2 * j + 1] += w * f.y; + } + } + } + if (t == 0) + { + sMax[localRow][g] = groupMax; + sSum[localRow][g] = sum; + } +#pragma unroll + for (int j = 0; j < 8; ++j) + { + sAcc[localRow][g][t * 8 + j] = acc[j]; + } + __syncthreads(); + if (g != 0 || !validRow) + { + return; + } + float rowMax = -CUDART_INF_F; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + rowMax = fmaxf(rowMax, sMax[localRow][k]); + } + float total = 0.f; + float o[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + if (sMax[localRow][k] == -CUDART_INF_F) + { + continue; + } + float const scale = exp2f(sMax[localRow][k] - rowMax); + total += scale * sSum[localRow][k]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + o[j] += scale * sAcc[localRow][k][t * 8 + j]; + } + } + float const inv = total > 0.f ? 1.f / total : 0.f; + int const token = row / p.headsPerKv; + int const head = row % p.headsPerKv; + __nv_bfloat162 packed[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + packed[j] = __floats2bfloat162_rn(o[2 * j] * inv, o[2 * j + 1] * inv); + } + size_t const outRow + = (static_cast(b) * p.seqLenQ + token) * (p.numKvHeads * p.headsPerKv) + h * p.headsPerKv + head; + reinterpret_cast(p.out + outRow * kHeadDim)[t] = *reinterpret_cast(packed); + if (p.q8 == nullptr) + { + return; + } + // MXFP8 copy of the BF16 row: a 32-column block is 4 neighbouring threads x 8 columns (lanes 0..15 of warp 0 + // hold the row). Block scale = the UE8M0 exponent of ceil(amax / 448), values = satfinite E4M3 of x * 2^-scale. + float x[8]; + float amax = 0.f; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = __bfloat1622float2(packed[j]); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + amax = fmaxf(amax, fmaxf(fabsf(f.x), fabsf(f.y))); + } + amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 2)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + uint16_t q2[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2[j]) : "f"(x[2 * j + 1] * qs), "f"(x[2 * j] * qs)); + } + int const tok = b * p.seqLenQ + token; + int const col = (h * p.headsPerKv + head) * kHeadDim + t * 8; + *reinterpret_cast(p.q8 + static_cast(tok) * (p.numSfKTiles * 128) + col) + = *reinterpret_cast(q2); + if (t % 4 == 0) + { + int const kSf = col / 32; + int64_t const off = static_cast(tok / 128) * p.numSfKTiles * 512 + (kSf / 4) * 512 + (tok % 32) * 16 + + ((tok % 128) / 32) * 4 + kSf % 4; + p.sf[off] = static_cast(e); + } +} + +} // namespace + +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) +{ + if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((params.rows + kRowsPerCta - 1) / kRowsPerCta, params.numKvHeads, params.batch); + config.blockDim = dim3(kThreads); + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + return params.bf16Partial ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) + : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); +} + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h new file mode 100644 index 000000000000..95af79e2b2ce --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h @@ -0,0 +1,77 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +//! The split-KV merge of the MiniMax-M3 TP2 full-context decode (PrimTS split-KV decode, separate-reduction scratch) +//! straight to BF16. See fullCtxReduce.cu. +namespace fullCtxReduce +{ + +constexpr int kHeadDim = 128; +constexpr int kMaxSplits = 128; + +struct Params +{ + void const* partialO; // [B, Hkv, S, R, 128] normalized 16-bit partial O per split + float const* partialStats; // [B, Hkv, S, R] log2-LSE per split + int const* seqLensKv; // [B] + __nv_bfloat16* out; // [B, SQ, Hkv * hR, 128] + int batch; + int numKvHeads; + int splits; // S, the planned split capacity (<= kMaxSplits) + int rows; // R = hR * SQ, rows [token, local head] + int seqLenQ; // SQ + int headsPerKv; // hR + int qTokensPerCta; // query tokens one producer CTA owns (grouped Q) + int tileKv; // producer K/V tile + int instsKv; // producer K/V instructions per loop step + bool causal; + bool bf16Partial; + uint8_t* q8 = nullptr; // optional E4M3 copy of out [B * SQ, Hkv * hR * 128] (the o_proj MXFP8 hand-off) + uint8_t* sf = nullptr; // its UE8M0 block scales, 128x4-swizzled as trtllm::mxfp8_quantize(x, True) lays them out + int numSfKTiles = 0; // Hkv * hR * 128 / 128 +}; + +//! Merges the active splits of every row in fp32 with one max rescale and writes BF16; with q8 set, also the MXFP8 +//! copy of the BF16 rows (bit for bit trtllm::mxfp8_quantize of out). With pdl the launch is +//! programmatic: the grid waits on the producer (griddepcontrol.wait) before reading its scratch. +//! +//! PDL contract: griddepcontrol.wait is the kernel's first instruction, so nothing (scratch, lengths, out) is read or +//! written before every prerequisite grid has completed and flushed; a producer's early trigger +//! (griddepcontrol.launch_dependents) only lets this grid launch and wait sooner. The kernel triggers its own +//! dependents right after its wait: the next kernel's griddepcontrol.wait still waits for this whole grid, so its +//! post-wait code sees all of out, and its pre-wait prologue must not read out. +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream); + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu new file mode 100644 index 000000000000..1942082c6df9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu @@ -0,0 +1,1264 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// SM107 (Rubin) TP4 copy of minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu for an sm_107 TP4 target: the row kernel at 4 +// ranks. Two changes in the arithmetic: the rank sum rounds to BF16 before the residual add, the rounding points of the +// stock all-reduce + fused residual RMSNorm on this target (bf16 all-reduce output, then bf16(residual + sum)); and +// every call but the MoE finalize runs the row kernel below, whose RMS statistic is the stock fused norm's reduction +// tree (bitwise the stock outputs). +// +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions): they are compiled for every sm_10x target of the build, the sm_100f family target +// included (which runs on SM107). A device pass for any other architecture compiles them to a trap, never to a silent +// empty body, and the host entry point checks that the kernel binary it launches holds sm_10x code (see +// checkBinary). +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace cg = cooperative_groups; + +namespace +{ + +constexpr int kEmpty = static_cast(0x80000000u); +constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k + +// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], +// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c +// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp +// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its +// arithmetic is the same source. (The TP2 target's mcAllreduceNormTp2.cu has the same epilogue.) +constexpr int kCapThreads = 256; +constexpr int kCapIters = 3; +constexpr int kCapWarps = kCapThreads / 32; +constexpr int kCapHidden = kCapThreads * kCapIters * 8; + +// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). +__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) +{ + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int q = 0; q < 8; ++q) + { + float const v = __bfloat162float(e[q]); + ss += v * v; + } + return ss; +} + +__device__ __forceinline__ float cap_warp_sum(float v) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. +__device__ __forceinline__ void cap_store( + float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) +{ + uint4 const wr = *reinterpret_cast(w); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int q = 0; q < 8; ++q) + oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); + *reinterpret_cast(cap) = o; +} + +__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) +{ + if ((w & 0xFFFFu) == 0x8000u) + w &= 0xFFFF0000u; + if ((w >> 16) == 0x8000u) + w &= 0x0000FFFFu; + return w; +} + +__device__ __forceinline__ float bf16_lo(uint32_t w) +{ + return __uint_as_float(w << 16); +} + +// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's +// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals +// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the +// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. +__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) +{ + float value; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); + uint32_t result; + asm("{\n" + " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" + " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" + " setp.le.f32 p_zero, %1, 0f00000000;\n" + " mov.b32 bits, %1;\n" + " shr.b32 exp_biased, bits, 23;\n" + " and.b32 exp_biased, exp_biased, 255;\n" + " and.b32 mantissa, bits, 0x7FFFFF;\n" + " setp.ne.u32 p_has_mant, mantissa, 0;\n" + " selp.u32 bump, 1, 0, p_has_mant;\n" + " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" + " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" + " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" + " @p_tiny_sub mov.u32 bump, 0;\n" + " add.u32 r, exp_biased, bump;\n" + " setp.gt.u32 p_ovf, r, 254;\n" + " selp.u32 r, 254, r, p_ovf;\n" + " selp.u32 %0, 0, r, p_zero;\n" + "}\n" + : "=r"(result) + : "f"(value)); + return result; +} + +__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) +{ + float inv; + asm("{\n" + " .reg .s32 new_exp;\n" + " .reg .b32 float_bits;\n" + " .reg .pred p_zero;\n" + " setp.eq.u32 p_zero, %1, 0;\n" + " sub.s32 new_exp, 254, %1;\n" + " max.s32 new_exp, new_exp, 0;\n" + " shl.b32 float_bits, new_exp, 23;\n" + " mov.b32 %0, float_bits;\n" + " @p_zero mov.b32 %0, 0;\n" + "}\n" + : "=f"(inv) + : "r"(ue8m0)); + return inv; +} + +// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. +__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) +{ + uint16_t pair; + asm("{\n" + " .reg .f32 f0, f1;\n" + " mul.f32 f0, %1, %3;\n" + " mul.f32 f1, %2, %3;\n" + " min.f32 f0, f0, 0f43E00000;\n" + " max.f32 f0, f0, 0fC3E00000;\n" + " min.f32 f1, f1, 0f43E00000;\n" + " max.f32 f1, f1, 0fC3E00000;\n" + " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" + "}\n" + : "=h"(pair) + : "f"(lo), "f"(hi), "f"(inv)); + return pair; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ float bf16_hi(uint32_t w) +{ + return __uint_as_float(w & 0xFFFF0000u); +} + +__device__ __forceinline__ int4 ld_volatile_v4(int const* p) +{ + int4 v; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +struct FinalizeArgs +{ + __nv_bfloat16 const* fc2; // [permuted rows, H] routed-expert FC2 output (null: plain input x) + int const* perm; // [M, topk] expanded index -> permuted row, -1 = dropped + void const* scales; // [M, topk] routing weights (fp32 or bf16) + int scale_fp32; + int topk; +}; + +// grid = min(M, MAXT) * CLUSTER CTAs, cluster (CLUSTER, 1, 1), THREADS threads. +// OPTS (each bit-identical): 1 residual + norm weight loaded at entry (before the poll); 2 cluster RMS exchange by +// st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated right after the poll; 8 max-smem +// carveout so the next kernel's CTAs co-reside early. +// OPTS 16 (NOT bit-identical with the others; a different boundary): the Eagle3 drafter's boundary, whose stock is a +// plain BF16 all-reduce, then FlashInfer fused_add_rmsnorm: the RMS statistic and the normalization take the +// unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded to BF16). +template +__global__ void __launch_bounds__(THREADS) mc_allreduce_norm_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, FinalizeArgs fin, int M, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int world, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ alignas(16) float cluster_sq[CLUSTER]; + extern __shared__ __align__(16) uint4 cap_slice[]; // [cols / 8] this CTA's updated chunks (cap_out only) + // Option 32, ONESHOT order (H 6144, CLUSTER 8, not the drafter option 16): every CTA's 3 active warp sums gathered + // in every CTA (global warp 3c + w), then reduced as the stock TP4 oneshot fusion does (see the consumer below). + constexpr bool kOneshotOrder = (OPTS & 32) != 0; + static_assert(!kOneshotOrder || (CLUSTER == 8 && THREADS == 128 && (OPTS & 16) == 0), "option 32: CLUSTER 8 x 128"); + constexpr int kAW = 3; + __shared__ alignas(16) float cluster_wsq[CLUSTER * kAW]; + __shared__ float oneshot_total; + __shared__ alignas(8) unsigned long long xbar; + cg::cluster_group cluster = cg::this_cluster(); + if constexpr ((OPTS & 2) != 0) + { + if (threadIdx.x == 0) + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(b) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + } + int const crank = static_cast(cluster.block_rank()); + int const t = blockIdx.x / CLUSTER; + int const cols = H / CLUSTER; + int const c0 = crank * cols + threadIdx.x * 8; + bool const active = t < M && threadIdx.x * 8 < cols; + int64_t const half_words = static_cast(MAXT) * world * (H / 2); + + // prepushed 1: the producer (m3_mx_gemv) already pushed this rank's rows into the exchange slots; 2: and it + // launched this grid only after its pushes (its grid wait passed), so the slots are polled without a grid wait; + // 3: as 2, with x in set A and x2 in set B. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const parity = *reinterpret_cast(flags + blockIdx.x) & 1; + uint4 rraw = make_uint4(0u, 0u, 0u, 0u); + uint4 wraw = make_uint4(0u, 0u, 0u, 0u); + if constexpr ((OPTS & 1) != 0) + { + // the residual row was written grids earlier (prepushed >= 2: the producer launched this grid only after its + // own grid wait; otherwise the grid wait above) + if (active) + { + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + wraw = *reinterpret_cast(norm_w + c0); + } + } + // Push this rank's slice into every rank's buffer. + if (active && !prepushed) + { + uint4 raw; + if (fin.fc2 != nullptr) + { + // MoE finalize: fp32 sum of the routing-weighted top-k expert rows plus the shared expert, rounded once. + float f[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + if (fin.topk <= kFinMaxTopk) + { + // All top-k indices and scales, then every expert-row load in flight at once, then the same in-order + // FP32 accumulation as the loop below (one memory round trip instead of topk; bit-identical). + int prow[kFinMaxTopk]; + float sc[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + prow[k] = -1; + sc[k] = 0.f; + if (k < fin.topk) + { + int const e = t * fin.topk + k; + prow[k] = fin.perm[e]; + sc[k] = fin.scale_fp32 ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + } + } + uint4 rv[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] >= 0) + rv[k] = *reinterpret_cast(fin.fc2 + static_cast(prow[k]) * H + c0); + } +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] < 0) + continue; + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv[k]); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc[k]; + f[2 * q + 1] += v.y * sc[k]; + } + } + } + else + for (int k = 0; k < fin.topk; ++k) + { + int const e = t * fin.topk + k; + int const prow = fin.perm[e]; + if (prow < 0) + continue; + float const sc = fin.scale_fp32 + ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + uint4 const rv = *reinterpret_cast(fin.fc2 + static_cast(prow) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc; + f[2 * q + 1] += v.y * sc; + } + } + if (x2 != nullptr) + { + // As stock: the routed output is rounded to BF16 first, then the shared expert is added (one more + // BF16 rounding of the FP32 sum). + uint4 const rv = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const routed = __bfloat1622float2(__floats2bfloat162_rn(f[2 * q], f[2 * q + 1])); + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] = routed.x + v.x; + f[2 * q + 1] = routed.y + v.y; + } + } + __nv_bfloat162* o = reinterpret_cast<__nv_bfloat162*>(&raw); +#pragma unroll + for (int q = 0; q < 4; ++q) + o[q] = __floats2bfloat162_rn(f[2 * q], f[2 * q + 1]); + } + else + { + raw = *reinterpret_cast(x + static_cast(t) * H + c0); + } + if (fin.fc2 == nullptr && x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * world + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float s[8] = {}; + float sq = 0.f; + if (active) + { + // The residual row (written grids earlier) and every rank's slot are loaded together; slots still empty are + // re-polled. Ranks are then summed in rank order, as before. + if constexpr ((OPTS & 1) == 0) + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + constexpr int kMaxWorld = 8; + int* const p0 = ws_uc + parity * half_words + static_cast(t) * world * (H / 2) + c0 / 2; + int4 v[kMaxWorld]; + int4 v2[kMaxWorld]; + bool const two = prepushed == 3; + int* const p1 = p0 + 2 * half_words; // set B + uint32_t pending = (1u << world) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (r < world) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr ((OPTS & 16) != 0) + { + float const x0 = s[2 * q] + rf.x, x1 = s[2 * q + 1] + rf.y; + s[2 * q] = x0; + s[2 * q + 1] = x1; + sq += x0 * x0 + x1 * x1; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + if constexpr (kOneshotOrder) + { + // stock oneshot: acc += v * v, one column at a time + sq += uf.x * uf.x; + sq += uf.y * uf.y; + } + else + sq += uf.x * uf.x + uf.y * uf.y; + } + } + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr ((OPTS & 16) == 0) + { + if (cap_out != nullptr) + cap_slice[threadIdx.x] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + if constexpr ((OPTS & 4) != 0) + { + // every CTA (active or not) rotates its call parity once its own poll is done (read only by later grids) + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffffu, sq, off); + if ((threadIdx.x & 31) == 0) + warp_sq[threadIdx.x / 32] = sq; + __syncthreads(); + if constexpr ((OPTS & 2) != 0) + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); // peers' xbar initialized (long since) + if (threadIdx.x == 0) + { + if constexpr (kOneshotOrder) + { + // Every CTA gets this CTA's kAW warp sums at [kAW * crank + w]. + if constexpr ((OPTS & 2) == 0) + { + for (int c = 0; c < CLUSTER; ++c) + for (int w = 0; w < kAW; ++w) + *cluster.map_shared_rank(&cluster_wsq[kAW * crank + w], c) = warp_sq[w]; + } + else + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * kAW * 4) + : "memory"); + for (int c = 0; c < CLUSTER; ++c) + { + uint32_t rb; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); + for (int w = 0; w < kAW; ++w) + { + uint32_t rd; + uint32_t const d + = static_cast(__cvta_generic_to_shared(&cluster_wsq[kAW * crank + w])); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), + "r"(__float_as_uint(warp_sq[w])), "r"(rb) + : "memory"); + } + } + } + } + else + { + float total = 0.f; +#pragma unroll + for (int w = 0; w < THREADS / 32; ++w) + total += warp_sq[w]; + if constexpr ((OPTS & 2) == 0) + for (int c = 0; c < CLUSTER; ++c) + *cluster.map_shared_rank(&cluster_sq[crank], c) = total; + else + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t const d = static_cast(__cvta_generic_to_shared(&cluster_sq[crank])); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * 4) + : "memory"); + for (int c = 0; c < CLUSTER; ++c) + { + uint32_t rd, rb; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), + "r"(__float_as_uint(total)), "r"(rb) + : "memory"); + } + } + } + } + if constexpr ((OPTS & 2) == 0) + cluster.sync(); + else + { + // every peer's sum landed (st.async complete_tx on this CTA's barrier); no cluster-wide release-arrive + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t done = 0; + while (!done) + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(b) + : "memory"); + } + if constexpr ((OPTS & 16) == 0) + { + if (cap_out != nullptr) + { + // The row's chunks are spread over the cluster (CTA c holds chunks [c * per, (c + 1) * per)); every CTA + // forms capture_norm's statistic itself, its 128 threads standing in for capture threads ct and ct + 128. + __shared__ float cap_red[kCapWarps]; + int const per = cols / 8; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + float ca = 0.f, cb = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + { + int const ka = it * kCapThreads + static_cast(threadIdx.x); + int const kb = ka + kCapThreads / 2; + ca = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[ka % per], ka / per), ca); + cb = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[kb % per], kb / per), cb); + } + // Every peer read is done before any CTA moves on (and later exits). + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + ca = cap_warp_sum(ca); + cb = cap_warp_sum(cb); + if ((threadIdx.x & 31) == 0) + { + cap_red[threadIdx.x / 32] = ca; + cap_red[threadIdx.x / 32 + kCapWarps / 2] = cb; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + if (active) + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + if constexpr (kOneshotOrder) + { + // The stock TP4 oneshot fusion at H 6144 (4 CTAs x 192 threads): blockReduceSumV2 per CTA = a 32-lane + // butterfly over its 6 warp sums (lanes >= 6 zero), lane 0's value; then thread 0 adds the 4 CTA sums in + // order from 0.f. Global warps 6g..6g+5 = this cluster's CTAs 2g, 2g+1. + if (threadIdx.x < 32) + { + int const lane = threadIdx.x; + float acc = 0.f; + for (int g = 0; g < 4; ++g) + { + float val = lane < 6 ? cluster_wsq[6 * g + lane] : 0.f; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + val += __shfl_xor_sync(0xffffffffu, val, off); + acc += __shfl_sync(0xffffffffu, val, 0); + } + if (lane == 0) + oneshot_total = acc; + } + __syncthreads(); + } + if (active) + { + float total = 0.f; + if constexpr (kOneshotOrder) + total = oneshot_total; + else + { +#pragma unroll + for (int c = 0; c < CLUSTER; ++c) + total += cluster_sq[c]; + } + float const inv = rsqrtf(total / static_cast(H) + eps); + if constexpr ((OPTS & 1) == 0) + wraw = *reinterpret_cast(norm_w + c0); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + } + if (out_q != nullptr) + { + // MXFP8 copy of the normed row for the next GEMM, as the CuTe DSL mxfp8 quantize of `normed` (the tree's + // decode quantize, MX-TAC; see mxfp8_ue8m0): four threads hold one 32-column scale block. + float amax = 0.f; + if (active) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + if (active) + { + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((threadIdx.x & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } + } + if constexpr ((OPTS & 4) == 0) + { + __syncthreads(); + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +// Refuses (raises) to launch a kernel whose loaded binary holds no sm_10x code (PTX and SASS version >= 100): the +// device passes for other architectures compile the kernels to a trap, and a build that left out every sm_10x target +// must not reach the GPU at all. +template +bool checkBinary(Kernel kernel, char const* name) +{ + cudaFuncAttributes attr{}; + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, reinterpret_cast(kernel))); + TLLM_CHECK_WITH_INFO(attr.ptxVersion >= 100 && attr.binaryVersion >= 100, + "%s: the loaded binary holds no sm_10x code (PTX %d, SASS %d): rebuild with an sm_100f / sm_107a target", name, + attr.ptxVersion, attr.binaryVersion); + return true; +} + +template +void launch(MinimaxM3Sm107Tp4McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) +{ + int const M = p.numTokens, H = p.hidden; + TLLM_CHECK_WITH_INFO(H % (CLUSTER * 8) == 0 && H / CLUSTER / 8 <= THREADS, "unsupported hidden size %d", H); + TLLM_CHECK_WITH_INFO(M >= 1 && M <= MAXT, "%d tokens do not fit the %d-token bucket", M, MAXT); + auto kernel = mc_allreduce_norm_kernel; + static bool const binaryOk = checkBinary(kernel, "mc_allreduce_norm_kernel"); + (void) binaryOk; + static bool const configured = [&]() + { + bool ok = CLUSTER <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess; + if ((OPTS & 8) != 0) + ok = ok + && cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared) + == cudaSuccess; + if (!ok) + cudaGetLastError(); + return ok; + }(); + TLLM_CHECK_WITH_INFO(configured, "cluster size not supported"); + FinalizeArgs const fin{p.fc2, p.perm, p.scales, p.scaleFp32 ? 1 : 0, p.topk}; + cudaLaunchConfig_t cfg = {}; + // The grid is sized to the call's rows: min(M, MAXT) clusters. Each (row, slice) slot keeps its own parity flag, + // read by the producers before they push and rotated only by the slot's CTA after it polled: a slot that a call + // does not use keeps its parity, and every rank runs the same calls, so the parities stay consistent when M + // changes between calls. + cfg.gridDim = dim3((M < MAXT ? M : MAXT) * CLUSTER); + cfg.blockDim = dim3(THREADS); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO((OPTS & 16) == 0 && H == kCapHidden && 2 * THREADS == kCapThreads, + "the capture epilogue takes H %d on the 128-thread cluster kernel, not the drafter boundary", kCapHidden); + cfg.dynamicSmemBytes = static_cast(H / CLUSTER / 8) * sizeof(uint4); + } + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = CLUSTER; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, fin, M, H, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, p.world, p.outNormed, p.outResidual, p.outQuant, p.outSf, (H / 32 + 3) / 4, + p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +// ws_uc / ws_mc are this rank's unicast and multicast views of ONE exchange buffer (they alias): no __restrict__. +// prepushed >= 2 contract: the producer (m3_mx_gemv's o_proj push, the m3_moe engine's finalize push) passes its own +// grid dependency, pushes its rows into the exchange, and only then releases this grid (griddepcontrol +// .launch_dependents); this kernel therefore skips griddepcontrol.wait and reads only the pushed (Lamport-polled) +// words, its parity flags and the residual / norm weight, which were written two or more grids earlier. +// +// The TP4 row kernel (4 ranks): one CTA of H / 8 = 768 threads per token row, used for every call but the MoE finalize. +// Its residual add + RMSNorm is the stock TP4 all-reduce fusion (allReduceFusionKernels.cu, oneshot Lamport), the fused +// norm after the all-reduce on this target, operation for operation at H 6144: thread j owns columns [8j, 8j + 8), its +// sum of squares is acc += v * v over them in order, then the warp butterfly 16..1, each group of 6 warps' sums through +// shared memory into a second 32-lane butterfly (lanes >= 6 zero), the 4 group sums added in order from 0.f (the stock +// oneshot kernel's 4-CTA cluster), rsqrtf(acc / H + eps) and (v * denom) * w. With the BF16 rank sum before the +// residual add, the normed rows, the updated residual and the MXFP8 copy are bitwise the stock ones. The exchange slots +// and the per-(row, slice) parity flags keep the cluster kernel's layout (slice = 96 threads), so producers that push +// (the o_proj GEMV) are unchanged. +constexpr int kSlices = 8; +constexpr int kSliceThreads = 96; +constexpr int kRowThreads = 768; +// The stock TP4 fused norm's launch at H 6144 (allreduce_fusion_kernel_launcher): 768 threads per token start as 8 +// CTAs of 96, then widen to CTAs of at least 128 threads: 4 CTAs of 192 per row cluster. +constexpr int kOneshotBlocks = 4; +constexpr int kOneshotThreads = kRowThreads / kOneshotBlocks; +static_assert(kOneshotThreads % 32 == 0 && kOneshotThreads / 32 <= 32, "one second-level butterfly per CTA"); + +// grid = M CTAs (one per row), THREADS = H / 8 threads. WORLD ranks (compile time). DRAFTER (opts bit 16): the Eagle3 +// drafter boundary, whose stock is a plain BF16 all-reduce then FlashInfer fused_add_rmsnorm: the RMS statistic and +// the normalization take the unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded). +// Its reduction is FlashInfer's CuTe-DSL FusedAddRMSNormKernel at H 6144 (what flashinfer.norm.fused_add_rmsnorm runs +// on SM 10.7 in this image: 128-thread CTAs, 64 threads per row, cluster 1): FI thread t sums its 96 columns +// 512 b + 8 t + e in (b, e) order from 0.f, a butterfly over its warp with offsets 1, 2, .., 16, the row's 2 warp sums +// through a second ascending butterfly (lanes >= 2 zero); rstd = rsqrt.approx(div.rn(sum, H) + eps), y = (h * rstd) * +// w. Bitwise to the stock boundary (oracle 777774: normed and residual, eager and graph, 4 ranks, 128 calls). CARVE +// (opts bit 64): the same kernel, its own function with the max-smem carveout preference set, so the next kernel's CTAs +// (the QKV GEMV's big staged tiles) can co-reside on the SM while it runs (the cluster kernel's option 8). ws_uc / +// ws_mc alias (one exchange buffer): no __restrict__. +template +__global__ void __launch_bounds__(THREADS, 1) mc_allreduce_norm_row_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int maxt, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ float group_sq[kOneshotBlocks]; + extern __shared__ __align__(16) uint4 cap_row[]; // [THREADS] the updated row's chunks (cap_out, !DRAFTER only) + // DRAFTER: the row's FP32 sums h, regrouped into FlashInfer's per-thread column sets. + __shared__ float h_row[DRAFTER ? THREADS * 8 : 1]; + __shared__ float total_sh; + int const t = blockIdx.x; + int const j = threadIdx.x; + int const c0 = j * 8; + int const slice = j / kSliceThreads; + int64_t const half_words = static_cast(maxt) * WORLD * (H / 2); + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); + int const parity = *reinterpret_cast(flags + t * kSlices + slice) & 1; + uint4 const rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + uint4 const wraw = *reinterpret_cast(norm_w + c0); + if (!prepushed) + { + uint4 raw = *reinterpret_cast(x + static_cast(t) * H + c0); + if (x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * WORLD + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } + asm volatile("griddepcontrol.launch_dependents;"); + float s[8] = {}; + float sq = 0.f; + { + int* const p0 = ws_uc + parity * half_words + static_cast(t) * WORLD * (H / 2) + c0 / 2; + int* const p1 = p0 + 2 * half_words; // set B + bool const two = prepushed == 3; + int4 v[WORLD]; + int4 v2[WORLD]; + uint32_t pending = (1u << WORLD) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr (DRAFTER) + { + s[2 * q] += rf.x; + s[2 * q + 1] += rf.y; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + } + } +#pragma unroll + for (int e = 0; e < 8; ++e) + sq += s[e] * s[e]; + if constexpr (DRAFTER) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + h_row[c0 + e] = s[e]; + } + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr (!DRAFTER) + { + if (cap_out != nullptr) + cap_row[j] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + // The stock TP4 all-reduce's fused RMSNorm (allReduceFusionKernels.cu, allreduce_fusion_kernel_oneshot_lamport at + // H 6144): a cluster of kOneshotBlocks CTAs x kOneshotThreads threads per row; each CTA reduces its warps with + // blockReduceSumV2 (warp butterfly 16..1, the CTA's warp sums into a second 32-lane butterfly, lanes past its warp + // count zero), then thread 0 sums the CTAs' totals in cluster-rank order from 0.f. + sq += __shfl_xor_sync(~0, sq, 16); + sq += __shfl_xor_sync(~0, sq, 8); + sq += __shfl_xor_sync(~0, sq, 4); + sq += __shfl_xor_sync(~0, sq, 2); + sq += __shfl_xor_sync(~0, sq, 1); + if ((j & 31) == 0) + warp_sq[j / 32] = sq; + __syncthreads(); + if constexpr (DRAFTER) + { + constexpr int kFiThreads = 64; // FlashInfer CuTe threads per row at H 6144 + constexpr int kFiBlocks = THREADS * 8 / (kFiThreads * 8); // 12 column blocks of 512 + if (j < kFiThreads) + { + float acc = 0.f; +#pragma unroll + for (int b = 0; b < kFiBlocks; ++b) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + { + float const v = h_row[b * kFiThreads * 8 + j * 8 + e]; + acc = __fmaf_rn(v, v, acc); // the DSL's h * h + reduce contracts to FMA + } + } +#pragma unroll + for (int off = 1; off < 32; off <<= 1) + acc += __shfl_xor_sync(~0, acc, off); + if ((j & 31) == 0) + group_sq[j / 32] = acc; + } + __syncthreads(); + if (j < 32) + { + float val = j < kFiThreads / 32 ? group_sq[j] : 0.f; +#pragma unroll + for (int off = 1; off < 32; off <<= 1) + val += __shfl_xor_sync(~0, val, off); + if (j == 0) + total_sh = val; + } + } + else if (j < kOneshotBlocks * 32) + { + constexpr int kWarpsPerBlock = kOneshotThreads / 32; + int const g = j / 32; + int const lane = j & 31; + float val = lane < kWarpsPerBlock ? warp_sq[g * kWarpsPerBlock + lane] : 0.f; + val += __shfl_xor_sync(~0, val, 16); + val += __shfl_xor_sync(~0, val, 8); + val += __shfl_xor_sync(~0, val, 4); + val += __shfl_xor_sync(~0, val, 2); + val += __shfl_xor_sync(~0, val, 1); + if (lane == 0) + group_sq[g] = val; + } + __syncthreads(); + if (!DRAFTER && j == 0) + { + float acc = 0.f; +#pragma unroll + for (int g = 0; g < kOneshotBlocks; ++g) + acc += group_sq[g]; + total_sh = acc; + } + // The parity of this row's slices rotates once every thread has read the exchange (later grids only). + if (j % kSliceThreads == 0) + flags[t * kSlices + slice] = parity ^ 1; + __syncthreads(); + if constexpr (!DRAFTER) + { + if (cap_out != nullptr) + { + // The whole row is in this CTA: threads 0..255 run capture_norm's chains over the staged chunks. + __shared__ float cap_red[kCapWarps]; + if (j < kCapThreads) + { + float css = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + css = cap_chunk_sq(cap_row[it * kCapThreads + j], css); + css = cap_warp_sum(css); + if ((j & 31) == 0) + cap_red[j / 32] = css; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + float inv; + if constexpr (DRAFTER) + { + float const mean = __fdiv_rn(total_sh, static_cast(H)); + asm("rsqrt.approx.f32 %0, %1;" : "=f"(inv) : "f"(__fadd_rn(mean, eps))); + } + else + inv = rsqrtf(total_sh / H + eps); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + if (out_q != nullptr) + { + float amax = 0.f; +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((j & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +template +auto rowKernel() +{ + auto kernel = mc_allreduce_norm_row_kernel<4, kRowThreads, DRAFTER, CARVE>; + static bool const ok = [&]() + { + checkBinary(kernel, "mc_allreduce_norm_row_kernel"); + if constexpr (CARVE) + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared)); + return true; + }(); + (void) ok; + return kernel; +} + +void launchRow(MinimaxM3Sm107Tp4McAllreduceNormParams const& p, int32_t maxTokens, bool drafter, bool carve, bool pdl, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(p.hidden == kRowThreads * 8 && p.world == 4, "the row kernel takes H 6144 at 4 ranks"); + TLLM_CHECK_WITH_INFO(p.numTokens >= 1 && p.numTokens <= maxTokens, "%d tokens do not fit the %d-token bucket", + p.numTokens, maxTokens); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(p.numTokens); + cfg.blockDim = dim3(kRowThreads); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO( + !drafter && p.hidden == kCapHidden, "the capture epilogue takes H %d, not the drafter", kCapHidden); + cfg.dynamicSmemBytes = static_cast(kRowThreads) * sizeof(uint4); + } + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + auto kernel = drafter ? (carve ? rowKernel() : rowKernel()) + : (carve ? rowKernel() : rowKernel()); + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, p.hidden, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, maxTokens, p.outNormed, p.outResidual, p.outQuant, p.outSf, + (p.hidden / 32 + 3) / 4, p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +constexpr int kCluster = 8; +constexpr int kThreads = 128; + +} // namespace + +bool isMinimaxM3Sm107Tp4McAllreduceNormSupported(int32_t maxTokens, int32_t opts) +{ + opts &= ~(64 | 128); // the row kernel's carveout instance; 128 only routes the call (invoke) + return (maxTokens == 8 && (opts == 7 || opts == 15 || opts == 23 || opts == 31 || opts == 39 || opts == 47)) + || ((maxTokens == 16 || maxTokens == 48) && (opts == 15 || opts == 47)); +} + +// Option 128 routes a world-4 call to the 8-CTA cluster kernel instead of the row kernel: the drafter boundary's +// cluster form (opts 23 / 31 = 16 | 7 / 16 | 15, dispatch.DRAFT_MC_AR_CLUSTER). Deterministic and exact-class to the +// row kernel's FlashInfer order: the updated residual is bitwise, the RMS sum of squares is the cluster's (per-CTA +// warps, then the 8-CTA total), so a normed element can differ by one BF16 rounding step. +void invokeMinimaxM3Sm107Tp4McAllreduceNorm(MinimaxM3Sm107Tp4McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream) +{ + if (params.fc2 == nullptr && params.world == 4 && params.hidden == kRowThreads * 8 && (opts & (32 | 128)) == 0) + { + launchRow(params, maxTokens, (opts & 16) != 0, (opts & 64) != 0, pdl, stream); + return; + } + opts &= ~(64 | 128); + if (maxTokens == 8 && opts == 39) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 47) + { + launch(params, pdl, stream); + } + else if (maxTokens == 16 && opts == 47) + { + launch(params, pdl, stream); + } + else if (maxTokens == 48 && opts == 47) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 7) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 23) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 31) + { + launch(params, pdl, stream); + } + else if (maxTokens == 16 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 48 && opts == 15) + { + launch(params, pdl, stream); + } + else + { + TLLM_THROW("unsupported (bucket %d, opts %d)", maxTokens, opts); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h new file mode 100644 index 000000000000..85938182c098 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h @@ -0,0 +1,83 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Operands of the MiniMax-M3 multicast all-reduce + residual + RMSNorm kernel (decode, at most 48 tokens). +//! +//! Per rank the kernel forms the partial `bf16(x + x2)` (or, in finalize mode, +//! `bf16(bf16(sum_k scale * fc2[perm]) + x2)`, the routed sum rounded before the shared expert is added, as stock), +//! pushes it through the multicast mapping `wsMc` into every rank's exchange, polls its own unicast view `wsUc`, +//! sums the ranks in FP32 in rank order, adds the residual (BF16 round) and applies RMSNorm. With `outQuant` it +//! also writes the MXFP8 copy of the normed rows (E4M3 [numTokens, hidden] and swizzled UE8M0 scales). +struct MinimaxM3Sm107Tp4McAllreduceNormParams +{ + __nv_bfloat16 const* x{nullptr}; //!< [numTokens, hidden] partial; in finalize mode the FC2 rows [rows, hidden]. + __nv_bfloat16 const* x2{nullptr}; //!< Optional second addend [numTokens, hidden]. + __nv_bfloat16 const* fc2{nullptr}; //!< Finalize mode: the permuted FC2 output (same buffer as x); null otherwise. + int32_t const* perm{nullptr}; //!< Finalize mode: [numTokens, topk] expanded index -> permuted row, -1 dropped. + void const* scales{nullptr}; //!< Finalize mode: [numTokens, topk] routing weights (fp32 or bf16). + bool scaleFp32{false}; + int32_t topk{0}; + int32_t numTokens{0}; + int32_t hidden{0}; + __nv_bfloat16 const* residual{nullptr}; //!< [numTokens, hidden] + __nv_bfloat16 const* normWeight{nullptr}; //!< [hidden] + float eps{0.F}; + int32_t* wsUc{nullptr}; //!< This rank's unicast view of the exchange (int32 words, re-armed after reading). + int32_t* wsMc{nullptr}; //!< The multicast mapping of the exchange. + int32_t* flags{nullptr}; //!< Per-CTA call-parity counters, [maxTokens * 8]. + int32_t rank{0}; + int32_t world{0}; + __nv_bfloat16* outNormed{nullptr}; + __nv_bfloat16* outResidual{nullptr}; + uint8_t* outQuant{nullptr}; //!< Optional E4M3 [numTokens, hidden]. + uint8_t* outSf{nullptr}; //!< Optional UE8M0 scales, 128x4-swizzled. + //! 0: push x (+ x2) here; 2: a producer already pushed x and launched this grid after its pushes; 3: as 2, with x + //! in exchange set A and x2 in set B. + int32_t prepushed{0}; + //! Optional Eagle3 capture of a capture layer's boundary: capOut[t * capStride + c] = the m3_spec_tail + //! capture_norm of the updated residual row with weight capWeight [hidden] and capEps, bit for bit (hidden 6144, + //! not the drafter boundary). + __nv_bfloat16* capOut{nullptr}; + int64_t capStride{0}; + __nv_bfloat16 const* capWeight{nullptr}; + float capEps{0.F}; +}; + +//! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). +bool isMinimaxM3Sm107Tp4McAllreduceNormSupported(int32_t maxTokens, int32_t opts); + +//! Launch with 8-CTA clusters of 128 threads, one cluster per token. `maxTokens` selects the exchange bucket; +//! `opts` the epilogue variant (every variant is bit-identical). +void invokeMinimaxM3Sm107Tp4McAllreduceNorm(MinimaxM3Sm107Tp4McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu new file mode 100644 index 000000000000..a0d5ac50025a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu @@ -0,0 +1,340 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h" + +#include +#include +#include + +// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP4 heads. +// +// The SM107 TP4 target's copy of the SM107 TP2 producerPw2.cu (itself a copy of minimaxM3/producerPw.cu, which the +// sm_103 targets keep), for its pre-wait step inputs. Same math and +// outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with --use_fast_math like +// the rest of the library), specialized to 16 query, 1 KV and 1 index head per rank (one index-K head), head_dim 128, +// rotary_dim 64, page 128, NVFP4 main KV and FP8 index-K. One warp per (token, packed head); K and V rows go to their +// head's block of the HND pools (the native kernel's head strides). +// +// The decode producer is latency bound: after griddepcontrol.wait the native kernel runs a chain of dependent global +// loads (packed row -> norm weight -> position id -> RoPE row -> cache slot -> quant scale). This kernel reads the +// constant norm weights and KV quant scale before griddepcontrol.wait and issues the packed row, position id and cache +// slot loads together after it, followed by the RoPE row; with `prewaitStepInputs` the position id, the slot and the +// RoPE row are read before the wait too, leaving only the packed row after it. With `early` it also calls +// griddepcontrol.launch_dependents on entry: consumers read its outputs only after their own griddepcontrol.wait, which +// waits for this grid to complete, so the early trigger only lets the next kernel launch and run its pre-wait prologue +// while this one runs. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace producerPw4 +{ +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotaryDim = 64; +constexpr int kPageSize = 128; +constexpr int kElems = kHeadDim / 32; +constexpr int kBlock = 256; + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ uint32_t sf_group_mask4() +{ + uint32_t laneId; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return 0xFu << (laneId & ~3u); +} + +__device__ __forceinline__ uint32_t fp32x8_to_e2m1(float (&a)[8]) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(a[0]), "f"(a[1]), "f"(a[2]), "f"(a[3]), "f"(a[4]), "f"(a[5]), "f"(a[6]), "f"(a[7])); + return val; +#else + return 0; +#endif +} + +__device__ __forceinline__ void store_fp8x4(__nv_fp8_e4m3* out, int64_t off, float const (&e)[kElems]) +{ + __nv_fp8x2_e4m3 const low(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const high(make_float2(e[2], e[3])); + *reinterpret_cast(out + off) = static_cast(low.__x) | (static_cast(high.__x) << 16); +} + +__device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(e[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = sf_group_mask4(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + *scaleOut = encodedScale.__x; + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float q[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + q[i] = __bfloat162float(rounded[i]) * outputScale; + return static_cast(fp32x8_to_e2m1(q)); +} + +template +__global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (kEarly) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * (kBlock / 32) + warpId; + int const tokenIdx = globalWarp / kTotalHeads; + int const localHead = globalWarp % kTotalHeads; + if (tokenIdx >= p.numTokens) + return; + + constexpr int kBegin = kHeadsQ; + constexpr int vBegin = kBegin + kHeadsKV; + constexpr int indexQBegin = vBegin + kHeadsKV; + constexpr int indexKHead = indexQBegin + kHeadsIndex; + bool const isQ = localHead < kBegin; + bool const isK = localHead >= kBegin && localHead < vBegin; + bool const isV = localHead >= vBegin && localHead < indexQBegin; + bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; + bool const isIndexK = localHead == indexKHead; + + // Constant data read before the wait: norm weights and the KV quant scale. + float w[kElems] = {}; + float quantScale = 0.0F; + __nv_bfloat16 const* normWeight + = isQ ? p.qWeight : (isK ? p.kWeight : (isIndexQ ? p.indexQWeight : p.indexKWeight)); + if (!isV) + { + uint2 const wv = *reinterpret_cast(normWeight + laneId * kElems); + float2 const w01 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.x)); + float2 const w23 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.y)); + w[0] = w01.x; + w[1] = w01.y; + w[2] = w23.x; + w[3] = w23.y; + } + if (isK || isV) + { + quantScale = p.kvQuantScale[isV ? 2 : 1]; + } + + int64_t const inputOffset = (static_cast(tokenIdx) * kTotalHeads + localHead) * kHeadDim + laneId * kElems; + bool const rotates = !isV && laneId * kElems < kRotaryDim; + bool const slotted = isK || isV || isIndexK; + int const coefficient = (laneId * kElems) % (kRotaryDim / 2); + // kPrewait: the position id, the cache slot and the RoPE row are read before the wait. They are per-step metadata + // written before the replay (host-to-device) or by plain-launched kernels earlier in the step (position_ids by the + // step's input preprocessing, the slots by msa_kv_update or its torch fallback), so they are final before this + // grid can start whatever its predecessors trigger; the RoPE table is constant. + int positionId = 0; + int slot = 0; + float4 c = make_float4(0.0F, 0.0F, 0.0F, 0.0F); + float4 s = c; + auto const loadStepInputs = [&]() + { + positionId = rotates ? p.positionIds[tokenIdx] : 0; + slot = slotted ? p.outCacheLoc[tokenIdx] : 0; + if (rotates) + { + int64_t const ropeRow = static_cast(positionId) * kRotaryDim; + c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); + s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); + } + }; + if constexpr (kPrewait) + { + loadStepInputs(); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + // The packed row comes from the predecessor grid; without kPrewait the step inputs are issued with it, so the + // native kernel's serial chain of dependent loads becomes two round trips. + uint2 const packed = *reinterpret_cast(p.packed + inputOffset); + if constexpr (!kPrewait) + { + loadStepInputs(); + } + float cosine[kElems] = {c.x, c.y, c.z, c.w}; + float sine[kElems] = {s.x, s.y, s.z, s.w}; + + float e[kElems]; + float sumSquares = 0.0F; +#pragma unroll + for (int pair = 0; pair < 2; ++pair) + { + float2 const v = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); + e[2 * pair] = v.x; + e[2 * pair + 1] = v.y; + if (!isV) + sumSquares += v.x * v.x + v.y * v.y; + } + + if (!isV) + { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + sumSquares += __shfl_xor_sync(0xffffffff, sumSquares, mask, 32); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kHeadDim) + p.eps); +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + constexpr int kPairOffset = (kRotaryDim / 2) / kElems; +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + int const dim = laneId * kElems + i; + float paired = __shfl_xor_sync(0xffffffff, e[i], kPairOffset); + if (dim < kRotaryDim) + { + if (dim < kRotaryDim / 2) + paired = -paired; + // The native kernel's contraction of e * cos + paired * sin; spelled out so this build matches it. + e[i] = __fmaf_rn(paired, sine[i], __fmul_rn(e[i], cosine[i])); + } + } + } + + if (isQ) + { + store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); + return; + } + if (isIndexQ) + { + store_fp8x4(p.indexQOut, + (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); + return; + } + + if (slot < 0) + return; + int const page = slot / kPageSize; + if (page >= p.numPages) + return; + int const withinPage = slot & (kPageSize - 1); + if (isIndexK) + { + store_fp8x4(p.indexK, + static_cast(page) * p.indexPageStride + static_cast(withinPage) * p.indexTokenStride + + laneId * kElems, + e); + return; + } + + int const head = isK ? localHead - kBegin : localHead - vBegin; + int const plane = isV ? 1 : 0; + int const scaleColumn = laneId / 4; + int64_t const scaleWithinPage = isV + ? static_cast(withinPage / 4) * 32 + static_cast(scaleColumn) * 4 + withinPage % 4 + : static_cast(withinPage) * 8 + scaleColumn; + int64_t const scaleOffset = static_cast(page) * p.scalePageStride + + static_cast(plane) * p.scalePlaneStride + static_cast(head) * p.scaleHeadStride + + scaleWithinPage; + uint8_t* scaleOut = laneId % 4 == 0 ? p.kvScale + scaleOffset : nullptr; + uint16_t const packedFp4 = quantize_fp4x4(e, quantScale, scaleOut); + int64_t const dataOffset = static_cast(page) * p.dataPageStride + + static_cast(plane) * p.dataPlaneStride + static_cast(head) * p.dataHeadStride + + static_cast(withinPage) * p.dataTokenStride + laneId * 2; + *reinterpret_cast(p.kvData + dataOffset) = packedFp4; +#endif +} + +template +cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) +{ + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; + int const warps = numTokens * kTotalHeads; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); + cfg.blockDim = dim3(kBlock); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); +} + +} // namespace + +cudaError_t invokeProducerPw4( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream) +{ + if (numTokens <= 0) + return cudaSuccess; + if (prewaitStepInputs) + { + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); + } + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); +} + +} // namespace producerPw4 +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h new file mode 100644 index 000000000000..6f8aa54a812b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP4 heads (16 query, 1 KV and +//! 1 index head per rank) with its constant loads before the PDL grid-dependency wait and an optional early dependents +//! trigger. Same outputs, bit for bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw4.cu. +namespace producerPw4 +{ + +struct Params +{ + __nv_bfloat16 const* packed; + __nv_fp8_e4m3* qOut; + __nv_fp8_e4m3* indexQOut; + uint8_t* kvData; + uint8_t* kvScale; + __nv_fp8_e4m3* indexK; + int const* outCacheLoc; + float const* kvQuantScale; + int64_t dataPageStride, dataPlaneStride, dataHeadStride, dataTokenStride; + int64_t scalePageStride, scalePlaneStride, scaleHeadStride; + int64_t indexPageStride, indexTokenStride, numPages; + int numTokens; + float eps; + __nv_bfloat16 const* qWeight; + __nv_bfloat16 const* kWeight; + __nv_bfloat16 const* indexQWeight; + __nv_bfloat16 const* indexKWeight; + float const* rotaryCosSin; + int const* positionIds; +}; + +//! The heads per rank this producer is built for. +constexpr int kHeadsQ = 16; +constexpr int kHeadsKV = 1; +constexpr int kHeadsIndex = 1; + +//! Launch over `numTokens` tokens, one warp per (token, packed head). `prewaitStepInputs` reads the position ids, the +//! cache slots and the RoPE rows before the grid-dependency wait (they must not be written by a PDL predecessor). +cudaError_t invokeProducerPw4( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream); + +} // namespace producerPw4 +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu new file mode 100644 index 000000000000..3d041421a92b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu @@ -0,0 +1,944 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h" + +#include "tensorrt_llm/common/assert.h" + +#include +#include + +// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). +// +// sm_107 (Rubin) TP4 copy of kernels/minimaxM3Sm107Tp2/sampledAccept.cu: identical algorithm and arithmetic as the +// shared kernels/minimaxM3/sampledAccept.cu, with the TP2 copy's launch (a row's CTAs as one thread-block cluster) and +// bounded record polls. The vocab shard limit stays the shared kernel's (kMaxVs 51200 >= the TP4 shard 50016). +// +// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +// active must stay on the generic path; the caller decides. +// +// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens +// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the +// nucleus with probability >= top_p, so most rows decide on their second exchange. +// +// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based +// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one +// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and +// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the +// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing +// and two call halves keep a rank one call ahead from overwriting a slot still being read. +// +// SA-1R (Args.one_round_k 9 / 82; MoE prototype moe/sa1r/6c753632ec, exact on 39,968 rows, eager + graph): each CTA +// also multicasts its largest bf16 logits with the round-0 record (3 per extra word). After round 0 every CTA knows M, +// the candidate logit cv and wc = __expf(cv - M); round 1 would accept iff S <= target, S = the fp32 sum (its tree) of +// w = __expf(__fmul_rn(x, tInv) - M) over every token with w > wc. x -> w is monotone (__fmul_rn by tInv > 0, fsub, +// __expf: checked exhaustively over all non-positive fp32 inputs), so if no CTA's unlisted value can beat wc (82: the +// bound tau; 9: the 9th listed), the listed w > wc are exactly round 1's summands, bitwise. Two fp32 sums of the same +// m non-negative terms differ by < m * 2^-23 * S, so accepting iff S_1r + m * 2^-21 * S_1r (round up) <= target +// implies round 1 accepts; m <= 1 is exact. Otherwise round 1 runs as before. An accepted row outputs the round-0 +// candidate, as the two-round kernel does, and consumes no Philox draw it would not. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace sampledAccept +{ +namespace +{ +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kMaxRec = kMaxWorld * kMaxSplits; +constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) +constexpr int kTagBits = 64 - 32 - kIdBits; +constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard +// Record polls trap after this many spins (the multicast all-reduce's bound): a missing writer becomes a CUDA error, +// never a hang. +constexpr long long kMaxSpins = 1LL << 31; + +// SA-1R variant kTopK: 0 none; K exact top-K list; 10 * K + P pooled (each thread keeps its top P plus its (P + 1)-th +// as a bound; the CTA lists the pool's top K plus tau = max(K-th listed, max thread bound)). +__host__ __device__ constexpr int mode_k(int m) +{ + return m >= 10 ? m / 10 : m; +} + +__host__ __device__ constexpr int mode_pool(int m) +{ + return m >= 10 ? m % 10 : 0; +} + +__host__ __device__ constexpr int mode_vals(int m) +{ + return mode_k(m) + (mode_pool(m) ? 1 : 0); +} + +// Words per record: key, logit of key, f2, f3, then the round-0 list (3 bf16 + a 16-bit tag per word). +template +struct Rec +{ + static constexpr int kXw = (mode_vals(kTopK) + 2) / 3; + static constexpr int kEx = 4 + kXw; +}; + +// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. +template +__host__ __device__ inline long long half_words(int world) +{ + return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * Rec::kEx; +} + +__device__ __forceinline__ uint32_t ordered(float v) +{ + uint32_t const b = __float_as_uint(v); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +// bf16 bits <-> an order-preserving 16-bit key. +__device__ __forceinline__ uint32_t ord16(uint32_t b) +{ + return (b & 0x8000u) ? (~b & 0xFFFFu) : (b | 0x8000u); +} + +__device__ __forceinline__ uint32_t unord16(uint32_t k) +{ + return (k & 0x8000u) ? (k & 0x7FFFu) : (~k & 0xFFFFu); +} + +// Philox4x32-10. +__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) +{ +#pragma unroll + for (int i = 0; i < 10; ++i) + { + uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; + uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; + c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); + k.x += 0x9E3779B9u; + k.y += 0xBB67AE85u; + } + return c; +} + +// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. +__device__ __forceinline__ float uniform(uint32_t bits) +{ + return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); +} + +// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. +__device__ __forceinline__ unsigned long long gkey(float lg, int id) +{ + return (static_cast(ordered(lg)) << kIdBits) + | static_cast((1u << kIdBits) - 1u - static_cast(id)); +} + +__device__ __forceinline__ int key_token(unsigned long long k) +{ + return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); +} + +__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) +{ + asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); +} + +__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) +{ + unsigned long long v; + asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); + return v; +} + +template +struct Shared +{ + __nv_bfloat16 row[kMaxVs / kSplits]; + float redf[kWarps]; + unsigned long long redk[kWarps]; + unsigned long long rk[kMaxRec]; + float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; + float wsum[kWarps]; + int owner; + int pick; + float pickx; + bool last; + float dS[kWarps]; // SA-1R: per-warp partial mass above the candidate + int dM[kWarps]; // SA-1R: per-warp count | incomplete << 16 + uint32_t wtop[kWarps][kTopK > 0 ? mode_k(kTopK) : 1]; // SA-1R: per-warp top-K keys + uint32_t wbound[kWarps]; // SA-1R pooled: per-warp max thread bound + uint32_t ctop[kTopK > 0 ? mode_vals(kTopK) : 1]; // SA-1R: the CTA's top-K keys (descending), by rank + unsigned long long rx[kMaxRec][Rec::kXw > 0 ? Rec::kXw : 1]; +}; + +template +__device__ __forceinline__ float block_max(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = s.redf[0]; + for (int i = 1; i < kWarps; ++i) + r = fmaxf(r, s.redf[i]); + __syncthreads(); + return r; +} + +// Fixed-order block sum: identical result for identical inputs. +template +__device__ __forceinline__ float block_sum(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = 0.f; + for (int i = 0; i < kWarps; ++i) + r += s.redf[i]; + __syncthreads(); + return r; +} + +template +__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + if ((threadIdx.x & 31) == 0) + s.redk[threadIdx.x >> 5] = v; + __syncthreads(); + unsigned long long r = s.redk[0]; + for (int i = 1; i < kWarps; ++i) + r = s.redk[i] > r ? s.redk[i] : r; + __syncthreads(); + return r; +} + +// Exclusive prefix of v over the threads in thread order, and the total. +template +__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + float inc = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + float const t = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) + inc += t; + } + if (lane == 31) + s.wsum[warp] = inc; + __syncthreads(); + float base = 0.f; + total = 0.f; + for (int i = 0; i < kWarps; ++i) + { + base += i < warp ? s.wsum[i] : 0.f; + total += s.wsum[i]; + } + __syncthreads(); + return base + inc - v; +} + +// Warp-wide integer max / sum / or (the redux instructions from sm_80; a butterfly otherwise, same results). +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_max_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + uint32_t const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_add_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_add_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_or_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_or_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v |= __shfl_xor_sync(0xffffffffu, v, o); + return v; +#endif +} + +// Insert key into a descending register list of K (fully unrolled: no local memory). +template +__device__ __forceinline__ void topk_insert(uint32_t (&a)[K], uint32_t key) +{ +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const hi = key > a[j] ? key : a[j]; + key = key > a[j] ? a[j] : key; + a[j] = hi; + } +} + +template +__device__ __forceinline__ void topk_pop(uint32_t (&a)[K]) +{ +#pragma unroll + for (int j = 0; j + 1 < K; ++j) + a[j] = a[j + 1]; + a[K - 1] = 0u; +} + +// Per-thread insertion with an early reject: after the first few elements almost every key fails the K-th test. +template +__device__ __forceinline__ void thread_insert(uint32_t (&a)[K], uint32_t key) +{ + if (key > a[K - 1]) + topk_insert(a, key); +} + +// The warp's K largest keys (descending) into s.wtop[warp]; published by the caller's next block barrier. +template +__device__ __forceinline__ void warp_topk(uint32_t (&a)[K], S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const m = warp_max_u32(a[0]); + if (a[0] == m && m != 0u) + topk_pop(a); + if (lane == 0) + s.wtop[warp][j] = m; + } +} + +// Pooled: K rounds over each lane's P pool keys (descending) and the warp's max thread bound. +template +__device__ __forceinline__ void warp_pool(uint32_t (&a)[P + 1], S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + uint32_t const bound = warp_max_u32(a[P]); + uint32_t pl[P]; +#pragma unroll + for (int j = 0; j < P; ++j) + pl[j] = a[j]; +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const m = warp_max_u32(pl[0]); + if (pl[0] == m && m != 0u) + topk_pop

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

(pl); + if (lane == 0) + s.wtop[warp][j] = m; + } + if (lane == 0) + s.wbound[warp] = bound; +} + +// A row's kSplits CTAs wait for each other's records, so they launch as one thread-block cluster: the hardware +// schedules a cluster's CTAs together, whatever other streams hold. Rows never wait on other rows. +template +__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) +{ + using R = Rec; + extern __shared__ __align__(16) unsigned char smem[]; + Shared& s = *reinterpret_cast*>(smem); + float constexpr kNegInf = -__builtin_huge_valf(); + int const tid = threadIdx.x; + int const row = blockIdx.x / kSplits; + int const split = blockIdx.x % kSplits; + int const world = a.world; + int const nrec = world * kSplits; + // The row's sampling parameters are step inputs (the spec metadata, written before the target forward), not the + // predecessor head GEMM's output, so with prewait_params they load before the dependency wait and overlap the + // head's tail; the same values either way. + float temperature = 0.f, topP = 0.f; + long long seed_in = 0, offset_in = 0; + auto load_params = [&]() + { + temperature = a.temperatures[row]; + topP = a.top_ps[row]; + seed_in = a.seeds[row]; + offset_in = a.offsets[row]; + }; + if (a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (!a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // early_trigger: the dependent grid (the first draft pass's input) may launch now and park at its own + // griddepcontrol.wait, which still returns only once this grid has completed and flushed; it skips the launch + // latency after this kernel ends. Otherwise dependents are released at the acceptance epilogue (the CTAs of + // splits 1.. exit without releasing). + if (a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const tempInv = 1.0f / (temperature + kTempEpsilon); + unsigned long long const seed = static_cast(seed_in); + unsigned long long const offset = static_cast(offset_in); + unsigned int const call = *reinterpret_cast(a.counters + 1); + // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of + // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round + // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten + // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. + auto ktag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; + }; + auto ftag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; + }; + // SA-1R round-0 list words: a 16-bit tag, never 0. + unsigned long long const xtag = (((call >> 1) << 1) & 0x7FFFu) | 0x8000u; + long long const rec_stride = static_cast(world) * kMaxSplits * R::kEx; + unsigned long long* const xu = a.ws_uc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride; + unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * R::kEx; + long long const round_stride = static_cast(kMaxRows) * rec_stride; + + // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- + __nv_bfloat16 const* src = a.logits + row * a.ld; + int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; + int const s0 = min(split * sl_len, a.Vs); + int const n = min(s0 + sl_len, a.Vs) - s0; + int const id0 = a.vocab_offset + s0; + float tmax = kNegInf; + // SA-1R: each thread's largest (ord16(bf16) << 16 | ~index) keys while staging (only rows with a nucleus cut). + bool const want_x = kTopK > 0 && topP < 1.0f; + constexpr int kK = mode_k(kTopK), kPool = mode_pool(kTopK), kVals = mode_vals(kTopK); + constexpr int kTk = kPool ? kPool + 1 : (kK > 0 ? kK : 1); + uint32_t tk[kTk]; +#pragma unroll + for (int j = 0; j < kTk; ++j) + tk[j] = 0u; + bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; + int const nvec = vec ? n / 8 : 0; + for (int v = tid; v < nvec; v += kThreads) + { + uint4 const w = *reinterpret_cast(src + s0 + 8 * v); + *reinterpret_cast(&s.row[8 * v]) = w; + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); +#pragma unroll + for (int j = 0; j < 8; ++j) + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); + if constexpr (kTopK > 0) + { + if (want_x) + { +#pragma unroll + for (int j = 0; j < 8; ++j) + { + uint32_t const key + = (ord16(__bfloat16_as_ushort(e[j])) << 16) | (0xFFFFu - static_cast(8 * v + j)); + if constexpr (kPool > 0) + topk_insert(tk, key); // branchless: P + 1 compare-swaps + else + thread_insert(tk, key); + } + } + } + } + for (int i = nvec * 8 + tid; i < n; i += kThreads) + { + s.row[i] = src[s0 + i]; + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); + if constexpr (kTopK > 0) + { + if (want_x) + topk_insert( + tk, (ord16(__bfloat16_as_ushort(src[s0 + i])) << 16) | (0xFFFFu - static_cast(i))); + } + } + if constexpr (kTopK > 0) + { + if (tid < kVals) + s.ctop[tid] = 0u; // published with s.wtop by block_max's barrier + if (want_x) + { + if constexpr (kPool > 0) + warp_pool(tk, s); // s.wtop / s.wbound published by block_max's barrier + else + warp_topk(tk, s); // s.wtop published by block_max's barrier + } + } + float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier + if constexpr (kTopK > 0) + { + // SA-1R: the CTA's top K of the warp lists by rank (keys are unique): a lane quad per key, no serial merge; + // s.ctop is read after the round-0 draw's first barrier + constexpr int kN = kWarps * kK; + constexpr int kSp = kN * 4 <= kThreads ? 4 : 1; // threads per key + if (want_x && tid < kN * kSp) + { + uint32_t const* const wl = &s.wtop[0][0]; + int const part = tid % kSp; + uint32_t const key = wl[tid / kSp]; + int rank = 0; +#pragma unroll + for (int j = part; j < kN; j += kSp) + rank += wl[j] > key ? 1 : 0; + if constexpr (kSp == 4) + { + rank += __shfl_xor_sync(0xffffffffu, rank, 1); + rank += __shfl_xor_sync(0xffffffffu, rank, 2); + } + if (part == 0 && key != 0u && rank < kK) + s.ctop[rank] = key; + } + } + float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z + unsigned long long ckey = 0ULL; + int const gid = a.rank * kMaxSplits + split; + for (int r = 0;; ++r) + { + // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) + float const ref = r == 0 ? smax : M; + float const wmin = r == 0 ? 0.f : wc; + float ts = 0.f; + for (int i = tid; i < n; i += kThreads) + { + float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); + ts += w > wmin ? w : 0.f; + } + float total; + float const excl = block_excl_scan(ts, total, s); + if constexpr (kTopK > 0) + { + // SA-1R round 0: the list words depend only on the staged slice, so they go out now (s.ctop / s.wbound are + // published by the barrier above), one lane per word, ahead of the main record: off the readers' path + if (r == 0 && want_x && tid < R::kXw) + { + uint32_t tau = 0u; + if constexpr (kPool > 0) + { + tau = s.ctop[kK - 1]; // the largest unlisted pool key is <= the K-th listed +#pragma unroll + for (int wq = 0; wq < kWarps; ++wq) + tau = s.wbound[wq] > tau ? s.wbound[wq] : tau; + } + unsigned long long w = xtag; +#pragma unroll + for (int t = 0; t < 3; ++t) + { + int const e = 3 * tid + t; + uint32_t const kv = e < kK ? s.ctop[e] : (kPool > 0 && e == kK ? tau : 0u); + // missing entries: -inf (bf16 0xFF80), w = 0 + uint32_t const b = kv != 0u ? unord16(kv >> 16) : 0xFF80u; + w |= static_cast(b) << (16 * (3 - t)); + } + st_sys(xm + 4 + tid, w); // round 0 = slot 0 + } + } + uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), + static_cast(offset), static_cast(offset >> 32)), + make_uint2(static_cast(seed), static_cast(seed >> 32))); + float const tgt = uniform(rnd.x) * total; + if (tid == 0) + s.owner = -1; + __syncthreads(); + if (ts > 0.f && excl <= tgt) + atomicMax(&s.owner, tid); + __syncthreads(); + if (tid == s.owner) + { + float acc = excl; + int pick = -1; + float px = 0.f; + if constexpr (kIlpWalk) + { + // SA-PW: the owner's strided tokens 8 at a time. The 8 x and w (the same expressions) are + // independent, so their MUFU work overlaps; the serial sum then adds them in the same order with + // the same early exit: the same acc sequence, pick and px as the one-token walk, bit for bit. + constexpr int kW = 8; + bool done = false; + for (int i0 = tid; i0 < n && !done; i0 += kW * kThreads) + { + float xs[kW], wv[kW]; +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + xs[k] = i < n ? __fmul_rn(__bfloat162float(s.row[i]), tempInv) : 0.f; + wv[k] = __expf(xs[k] - ref); + } +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + if (!done && i < n && wv[k] > wmin) + { + acc += wv[k]; + pick = i; + px = xs[k]; + if (acc > tgt) + done = true; + } + } + } + } + else + { + for (int i = tid; i < n; i += kThreads) + { + float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); + float const w = __expf(x - ref); + if (w > wmin) + { + acc += w; + pick = i; + px = x; + if (acc > tgt) + break; + } + } + } + s.pick = pick; + s.pickx = px; + } + __syncthreads(); + unsigned long long sb = 0ULL; + if (total > 0.f && s.owner >= 0) + sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); + float const smab = total; + unsigned long long const ktag = ktag_of(r); + unsigned long long const ftag = ftag_of(r); + unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; + bool const with_x = want_x && r == 0; + if (tid == 0) + { + unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; + st_sys(xw + 0, (sb << kTagBits) | ktag); + st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); + st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); + st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); + } + // ---- every record of the row, every rank --------------------------------------------------------------- + int const nw = with_x ? R::kEx : 4; + for (int t = tid; t < nrec * nw; t += kThreads) + { + int const q = t / nw, f = t % nw; + unsigned long long const* const p = xr + (q / kSplits * kMaxSplits + q % kSplits) * R::kEx + f; + unsigned long long v; + long long spins = 0; + if (f == 0) + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); + s.rk[q] = v >> kTagBits; + } + else if (f < 4) + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFFFFF00000000ULL) != ftag); + float const fv = __uint_as_float(static_cast(v)); + (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; + } + else + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFULL) != xtag); + s.rx[q][f - 4] = v; + } + } + __syncthreads(); + // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- + unsigned long long gk = 0ULL; + float gv = 0.f; + for (int q = 0; q < nrec; ++q) + { + if (s.rk[q] > gk) + { + gk = s.rk[q]; + gv = s.rv[q]; + } + } + bool accept; + if (r == 0) + { + M = kNegInf; + for (int q = 0; q < nrec; ++q) + M = fmaxf(M, s.r2[q]); + float Z = 0.f; + for (int q = 0; q < nrec; ++q) + Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); + target = topP * Z; + accept = !(topP < 1.0f); // no nucleus cut: the first draw stands + ckey = gk; + cv = gv; + if constexpr (kTopK > 0) + { + // SA-1R: every thread takes (record, entry) items; the per-warp partials are published by the round's + // barrier below and every thread combines them in warp order after it (no new barrier) + if (with_x) + { + float const wcand = __expf(cv - M); // = round 1's wmin + float part = 0.f; + int cnt = 0, inc = 0; + for (int t = tid; t < nrec * kVals; t += kThreads) + { + int const q = t / kVals, e = t % kVals; + uint32_t const b = static_cast(s.rx[q][e / 3] >> (16 * (3 - e % 3))) & 0xFFFFu; + float const w = __expf( + __fmul_rn(__bfloat162float(__ushort_as_bfloat16(static_cast(b))), tempInv) + - M); // w bitwise as round 1 computes it + if (w > wcand) + { + if (kPool > 0 && e == kK) + inc = 1; // tau (the largest unlisted value) is heavier than the candidate + else + { + part += w; + ++cnt; + if (kPool == 0 && e == kK - 1) + inc = 1; // the K-th is heavier: unlisted heavier tokens may exist + } + } + } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + part += __shfl_xor_sync(0xffffffffu, part, o); // any order: the margin covers it + cnt = static_cast(warp_add_u32(static_cast(cnt))); + inc = static_cast(warp_or_u32(static_cast(inc))); + if ((tid & 31) == 0) + { + s.dS[tid >> 5] = part; + s.dM[tid >> 5] = cnt | (inc << 16); + } + } + } + } + else + { + float S = 0.f; // mass strictly above the current candidate + for (int q = 0; q < nrec; ++q) + S += s.r2[q]; + accept = S <= target || gk == 0ULL; + if (!accept) + { + ckey = gk; + cv = gv; + } + } + wc = __expf(cv - M); + __syncthreads(); // records consumed before the next round's writes to s + bool decided = false; + if constexpr (kTopK > 0) + { + if (r == 0 && with_x) + { + float S1 = 0.f; + int m = 0, inc = 0; + for (int wq = 0; wq < kWarps; ++wq) + { + S1 += s.dS[wq]; + m += s.dM[wq] & 0xFFFF; + inc |= s.dM[wq] >> 16; + } + if (!inc) + { + float const margin = m <= 1 ? 0.f : S1 * static_cast(m) * (1.0f / 2097152.0f); + decided = __fadd_ru(S1, margin) <= target; + accept = accept || decided; + } + } + } + if (accept) + { + if (split != 0) + return; + if (tid == 0) + { + a.tokens[row] = key_token(ckey); + if (a.extra_attempts != nullptr && r > 1) + atomicAdd(a.extra_attempts, r - 1); + if (a.decided != nullptr && decided) + atomicAdd(a.decided, 1); + } + break; + } + } + + // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (!a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (tid == 0) + s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; + __syncthreads(); + if (!s.last) + return; + __threadfence(); + if (tid == 0) + { + a.counters[0] = 0u; + a.counters[1] = call + 1u; + } + if (tid < a.nreq) + { + int const* tok = a.tokens; + int const b = tid; + int const k1 = a.draft_len + 1; + int nacc = 1; + bool run = true; + for (int k = 0; k < a.draft_len; ++k) + { + int const tt = *reinterpret_cast(tok + b * k1 + k); + a.acc_tokens[b * k1 + k] = tt; + run = run && a.draft[b * a.draft_len + k] == tt; + nacc += run ? 1 : 0; + } + a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); + if (a.force_vec != nullptr) + nacc = a.force_vec[b]; + else if (a.force_total > 0) + nacc = a.force_total; + a.num_acc[b] = nacc; + a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; + if (a.next_new != nullptr) + a.next_new[b * a.next_stride] + = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); + } +} + +template +cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) +{ + static int max_clusters = 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(a.rows * kSplits); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Shared); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + if (max_clusters == 0) + { + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } + } + TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); + return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); +} + +template +cudaError_t launchRows(Args const& a, bool use_pdl, cudaStream_t stream) +{ + // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); + // more rows: 4 CTAs per row. + if (a.ilp_walk) + { + return a.rows <= 32 ? launch<8, kTopK, true>(a, use_pdl, stream) : launch<4, kTopK, true>(a, use_pdl, stream); + } + return a.rows <= 32 ? launch<8, kTopK, false>(a, use_pdl, stream) : launch<4, kTopK, false>(a, use_pdl, stream); +} + +} // namespace + +long long sampledAcceptWsWords(int world, int oneRoundK) +{ + TLLM_CHECK_WITH_INFO( + oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "SA-1R variant %d is not built", oneRoundK); + return 2 * (oneRoundK == 9 ? half_words<9>(world) : oneRoundK == 82 ? half_words<82>(world) : half_words<0>(world)); +} + +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) +{ + switch (a.one_round_k) + { + case 0: return launchRows<0>(a, usePdl, stream); + case 9: return launchRows<9>(a, usePdl, stream); + case 82: return launchRows<82>(a, usePdl, stream); + default: TLLM_CHECK_WITH_INFO(false, "SA-1R variant %d is not built", a.one_round_k); + } + return cudaErrorInvalidValue; +} + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h new file mode 100644 index 000000000000..aa52cf4a0559 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h @@ -0,0 +1,87 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ +//! Fused sampled (temperature + top-p) verification of the MiniMax-M3 Eagle3 target head over the TP-sharded +//! vocabulary, decode rows <= kMaxRows; see sampledAccept.cu. +namespace sampledAccept +{ + +constexpr int kMaxSplits = 8; // CTAs per verify row (at most) +constexpr int kMaxRows = 48; // verify rows per call +constexpr int kMaxWorld = 8; // TP ranks +constexpr int kMaxVs = 100352; // vocab shard width (TP2: 200064 / 2 = 100032, padded to a multiple of 512) +constexpr int kIdBits = 18; // token ids < 2^18 + +struct Args +{ + __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard + long long ld; + int rows, Vs, vocab_offset; + float const* temperatures; + float const* top_ps; + long long const* seeds; + long long const* offsets; + unsigned long long* ws_uc; + unsigned long long* ws_mc; + unsigned int* counters; // [0] finished rows, [1] call + int* tokens; // [rows] sampled token per row + int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows + int rank, world; + int const* draft; + int const* batch_idx; + int* acc_tokens; + int* num_acc; + long long* gather_ids; + int draft_len, nreq, force_total; + int* next_new; + int next_stride; + int const* force_vec; + int prewait_params; // nonzero: the row's sampling parameters load before the PDL dependency wait + int early_trigger; // nonzero: dependents are released right after the dependency wait, not at the epilogue + // SA-1R (one-round decision): 0 the two-round kernel; 9 each CTA's top-9 logits ride its round-0 record; 82 pooled + // (each thread's top 2, the CTA lists the pool's top 8 plus the bound of everything unlisted). A row whose first + // candidate is provably in the nucleus then decides in round 0, with the same token as the two-round kernel. + int one_round_k; + int* decided; // optional [1]: rows decided in round 0 (SA-1R), summed + int ilp_walk; // nonzero: the owner walk takes its tokens 8 at a time (same weights, same summation order) +}; + +//! Words (uint64) of the multicast exchange workspace for `world` ranks and SA-1R variant `oneRoundK`: two call halves. +long long sampledAcceptWsWords(int world, int oneRoundK = 0); + +//! Launch the sampled verification over a.rows * (8 or 4) co-resident CTAs (8 per row up to 32 rows). +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream); + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu new file mode 100644 index 000000000000..8cbf5d90175c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu @@ -0,0 +1,89 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kThreads = 128; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves). The +// instruction sequence of the Triton kernel's PTX at its defaults (host/silu/silu_and_mul_plain.ptx): g * -log2(e) as +// one product, ex2.approx, + 1, div.full 1 / that, g * s, (u + 0.0) as an IEEE add, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t siluMulBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 gl, gh, ul, uh, t0, t1, e0, e1, s0, s1, p0, p1, q0, q1, r0, r1, one, zero;\n\t" + ".reg .b16 hl, hh;\n\t" + "mov.b32 {hl, hh}, %1;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mul.rn.f32 t0, gl, 0fBFB8AA3B;\n\tmul.rn.f32 t1, gh, 0fBFB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 p0, gl, s0;\n\tmul.rn.f32 p1, gh, s1;\n\t" + "mov.b32 zero, 0f00000000;\n\tadd.rn.f32 q0, ul, zero;\n\tadd.rn.f32 q1, uh, zero;\n\t" + "mul.rn.f32 r0, q0, p0;\n\tmul.rn.f32 r1, q1, p1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +// Eight outputs per thread: 16-byte loads of the gate and the up half, one 16-byte store. +__global__ void __launch_bounds__(kThreads) + siluAndMulKernel(uint4 const* __restrict__ x, uint4* __restrict__ out, int64_t rows, int64_t vecs) +{ + int64_t const row = blockIdx.y; + int64_t const v = static_cast(blockIdx.x) * kThreads + threadIdx.x; + if (row >= rows || v >= vecs) + { + return; + } + uint4 const g = x[row * 2 * vecs + v]; + uint4 const u = x[row * 2 * vecs + vecs + v]; + uint4 o; + o.x = siluMulBf16x2(g.x, u.x); + o.y = siluMulBf16x2(g.y, u.y); + o.z = siluMulBf16x2(g.z, u.z); + o.w = siluMulBf16x2(g.w, u.w); + out[row * vecs + v] = o; +} + +} // namespace + +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream) +{ + int64_t const vecs = d / 8; + dim3 const grid(static_cast((vecs + kThreads - 1) / kThreads), static_cast(rows)); + siluAndMulKernel<<>>( + reinterpret_cast(x), reinterpret_cast(out), rows, vecs); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h new file mode 100644 index 000000000000..669e33eb5cf5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h @@ -0,0 +1,40 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +//! out[m, j] = bf16(silu(g) * (u + 0)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both +//! contiguous, d a multiple of 8, 16-byte aligned rows): bit for bit the Triton silu_and_mul kernel at its defaults +//! (no limit, alpha 1, beta 0), the plain SwiGLU of the Eagle3 drafter MLP. +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu new file mode 100644 index 000000000000..d51138cc15d2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu @@ -0,0 +1,746 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// MiniMax-M3 TP8 long-context sparse decode attention in one launch (8 query heads on one KV head). +// +// Grid: one cluster of 16 CTAs per query token; CTA `r` attends the token's `r`-th selected block. The QK and PV +// MMAs are transposed (S^T = K q^T, O^T = V^T P^T) so the 8 heads fill the MMA n dimension. Each CTA pushes its +// normalized partial over DSMEM so that CTA `r` holds (head r / 2, dims (r % 2) * 64 + [0, 64)) of all 16 partials, +// merges them, and writes the BF16 row and, optionally, its MXFP8 copy and swizzled UE8M0 scales. +// +// Per-block math: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and PV on FP16 tensor cores with +// FP32 accumulation, log2-domain softmax with either the linear denominator or the log-domain recurrence. +// +// PREWAIT stages everything that does not depend on the top-16 table before the grid dependency wait: the request's +// seq_lens entry, the q rows and the request's block-table row (in dynamic shared memory after Smem8, when it has at +// most kBtCap entries). The arithmetic is the same in both variants. +// +// ONE_DIV merge (MSA_ONE_DIV_MERGE): each CTA pushes its unnormalized partial with its row max m and sum l, and the +// merge forms o = sum_c w_c acc_c / sum_c w_c l_c with w_c = 2^(m_c - max m), one division per output, instead of +// normalizing every partial and weighting by its LSE. Deterministic (fixed chunk order, no atomics); a different +// rounding sequence from the Triton merge, inside the same band of the float64 reference. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kDim = 128; +constexpr int kPage = 128; +constexpr int kTopk = 16; +constexpr int kThreads = 128; +constexpr int kPitch = kDim + 8; // halves per smem row: 272 B, conflict-free ldmatrix +constexpr float kNegInf = -__builtin_huge_valf(); +constexpr int kHeads8 = 8; +constexpr int kHalfDim = kDim / 2; +constexpr int kBtCap = 8192; // block-table entries staged in shared memory (8192 x 128 = 1M tokens) + +using Params = MinimaxM3Tp8SparseDecodeLc8Params; + +struct Smem8 +{ + __half qs[kHeads8 * kPitch]; + __half vs[kPage * kPitch]; + __half ps[kHeads8 * kPitch]; // P^T: [head][token] + float recv[kTopk][kHalfDim]; // (head rank / 2, dims (rank % 2) * 64 + [0, 64)) partials of the 16 CTAs + float recv_lse[kTopk]; // the chunk LSEs (log2); with ONE_DIV, the chunk row maxima + float recv_l[kTopk]; // ONE_DIV: the chunk row sums + float part[kHalfDim]; + float ex[kTopk]; // the 16 chunk exponentials (ONE_DIV: weighted row sums), formed once per CTA + float wgt[kTopk]; // the 16 merge weights, formed once per CTA + float red[4][kHeads8]; +}; + +template +constexpr size_t smemBytes() +{ + return PREWAIT ? sizeof(Smem8) + kBtCap * sizeof(int) : sizeof(Smem8); +} + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t e2m1x2_to_f16x2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 s;\n cvt.u8.u32 s, %1;\n cvt.rn.f16x2.e2m1x2 %0, s;\n}" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t e4m3x2_to_f16x2(uint16_t pair) +{ + uint32_t out; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(pair)); + return out; +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ void ldsm_x4(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldsm_x4_t(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma(float (&c)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ void cluster_sync() +{ + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +__device__ __forceinline__ uint32_t map_cluster(void const* local, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smem_u32(local)), "r"(rank)); + return remote; +} + +__device__ __forceinline__ void st_cluster_f32(uint32_t addr, float a) +{ + asm volatile("st.shared::cluster.f32 [%0], %1;" ::"r"(addr), "f"(a) : "memory"); +} + +// One packed NVFP4 row (64 bytes, 128 values) with its 8 E4M3 scales -> 128 fp16 values in smem. +__device__ __forceinline__ void dequant_row(__half* dst, uint8_t const* packed, uint8_t const (&sc)[8], bool valid) +{ + uint4 const* src = reinterpret_cast(packed); + uint32_t* d = reinterpret_cast(dst); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + uint32_t words[4] = {w.x, w.y, w.z, w.w}; + // 16 bytes = 32 values = scale columns 2c and 2c + 1. +#pragma unroll + for (int h = 0; h < 2; ++h) + { + uint8_t const s = valid ? sc[2 * c + h] : 0; + uint32_t const sh = e4m3x2_to_f16x2(static_cast(s | (s << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); +#pragma unroll + for (int j = 0; j < 2; ++j) + { + uint32_t const word = words[2 * h + j]; +#pragma unroll + for (int b = 0; b < 4; ++b) + { + uint32_t const v2 = e2m1x2_to_f16x2((word >> (8 * b)) & 0xFF); + __half2 const x = __hmul2(*reinterpret_cast<__half2 const*>(&v2), scale); + d[c * 16 + h * 8 + j * 4 + b] = *reinterpret_cast(&x); + } + } + } + } +} + +// q: 8 heads x 128 e4m3 -> fp16; threads 0..63 = (head, 16-value chunk). +__device__ __forceinline__ void load_q(Smem8& s, Params const& p, int tok, int tid) +{ + if (tid < kHeads8 * 8) + { + int const h = tid >> 3; + int const c = (tid & 7) * 16; + uint4 const w = *reinterpret_cast(p.q + tok * p.q_tok + h * p.q_head + c); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; + uint32_t* d = reinterpret_cast(s.qs + h * kPitch + c); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + d[2 * i] = e4m3x2_to_f16x2(static_cast(words[i] & 0xFFFF)); + d[2 * i + 1] = e4m3x2_to_f16x2(static_cast(words[i] >> 16)); + } + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +// This CTA's 16-byte-aligned slice of [base, base + bytes) into L2, at most 16 KiB per bulk prefetch. +__device__ __forceinline__ void prefetch_l2_slice(uint8_t const* base, int64_t bytes) +{ + if (base == nullptr || bytes <= 0) + return; + int64_t const n = gridDim.x; + int64_t const slice = ((bytes + n - 1) / n + 15) & ~int64_t{15}; + int64_t const off = static_cast(blockIdx.x) * slice; + if (off >= bytes) + return; + int64_t const len = min(slice, bytes - off) & ~int64_t{15}; + for (int64_t o = 0; o < len; o += 16384) + { + uint32_t const sz = static_cast(min(int64_t{16384}, len - o)); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(base + off + o), "r"(sz) : "memory"); + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +template +__global__ void __launch_bounds__(kThreads) m3_tp8_lc8_attn_kernel(__grid_constant__ Params const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(128) uint8_t smem_raw[]; + Smem8& s = *reinterpret_cast(smem_raw); + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const g = lane >> 2; + int const qd = lane & 3; + int const tok = blockIdx.x / kTopk; + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + // The o_proj weights are read-only for the whole step: start pulling them into L2 now, while HBM is idle in the + // latency-bound attention phase (one lane of warps 0 and 1; nothing waits on these). + if (tid == 0) + prefetch_l2_slice(p.pf0, p.pf0_bytes); + else if (tid == 32) + prefetch_l2_slice(p.pf1, p.pf1_bytes); + if constexpr (!PREWAIT) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + + int const req = tok / p.dql; + int const qoff = tok - req * p.dql; + int const seq = p.seq_lens[req]; + int const kv_len = max(seq - p.dql + qoff + 1, 0); + int const nblocks = (kv_len + kPage - 1) / kPage; + int const real_topk = min(kTopk, nblocks); + // The next sparse layer's block that its CTA (tok, rank) selected at the previous step (a decode step selects + // mostly the blocks it selected one step earlier): its K/V pages and scales into L2 for that layer, from one lane + // of warp 2. The id buffer was written by a kernel of the previous step, long complete; a recorded id at or past + // this request's block count (another request's row) or a page outside the pool is skipped. + if (p.n1_next != nullptr && tid == 64) + { + int const nb = p.n1_next[tok * kTopk + rank]; + if (nb >= 0 && nb < nblocks) + { + int64_t const pg = p.block_table[req * p.bt_stride + nb]; + if (pg >= 0 && pg < p.n1_pages) + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_k + pg * p.k_blk), "r"(kPage * 64) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_v + pg * p.v_blk), "r"(kPage * 64) + : "memory"); + asm volatile( + "cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_ksc + pg * p.ksc_blk), "r"(kPage * 8) + : "memory"); + asm volatile( + "cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_vsc + pg * p.vsc_blk), "r"(kPage * 8) + : "memory"); + } + } + } + + // The request's block-table row, staged before the grid dependency (written before the graph runs). + int* const s_bt = reinterpret_cast(smem_raw + sizeof(Smem8)); + bool bt_smem = false; + if constexpr (PREWAIT) + { + int const nbt = (max(seq, 0) + kPage - 1) / kPage; + bt_smem = nbt <= kBtCap; + if (bt_smem) + { + for (int i = tid; i < nbt; i += kThreads) + { + s_bt[i] = __ldg(p.block_table + req * p.bt_stride + i); + } + } + // Before the wait only behind the select: an identity step has no scorer / select, so this grid's + // predecessor is the producer itself, which releases its dependents at entry, before it writes q. There q is + // read after the wait, whether the blocks come from the lengths (identity) or from the staged table. + if (!p.identity && !p.identity_step) + load_q(s, p, tok, tid); + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int blk = static_cast(rank) >= real_topk ? -1 + : p.identity ? static_cast(rank) + : p.topk[tok * p.t_tok + rank * p.t_k]; + // Fail closed on an id past the row (a stale or corrupt table), as the stock Triton producer does: the block is + // skipped. Read through the block table it would be another request's page, every token masked, and its softmax + // ex2(-inf - (-inf)) = NaN. + if (blk >= nblocks) + blk = -1; + bool const active = blk >= 0; + // Record this step's selection for this layer's next-step prefetch (read by the previous sparse layer's CTAs). + if (p.n1_self != nullptr && tid == 0) + p.n1_self[tok * kTopk + rank] = blk; + if constexpr (PREWAIT) + { + if (p.identity || p.identity_step) + load_q(s, p, tok, tid); + __syncthreads(); // s_bt and the q rows are complete + } + else + { + load_q(s, p, tok, tid); + } + + // Lane holds heads 2 * qd + {0, 1}; acc[mt]: dims warp * 32 + mt * 16 + g + {0, 8}. + float m_row[2] = {kNegInf, kNegInf}; + float l_row[2] = {0.F, 0.F}; + float acc[2][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + acc[j][i] = 0.F; + + // K rows n = warp * 32 + j * 8 + g stay packed in registers. As the MMA A operand, m-tile mt takes rows + // j = 2 mt (row g) and 2 mt + 1 (row g + 8); k-step ks takes packed bytes 8 ks + qd and 8 ks + 4 + qd. + uint32_t kw[4][16]; + uint2 ksc[4]; + if (active) + { + int64_t const page = bt_smem ? s_bt[blk] : p.block_table[req * p.bt_stride + blk]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const n = warp * 32 + j * 8 + g; + bool const valid = blk * kPage + n < kv_len; + uint4 const* src = reinterpret_cast(p.k + page * p.k_blk + n * 64); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 const w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + kw[j][4 * c] = w.x; + kw[j][4 * c + 1] = w.y; + kw[j][4 * c + 2] = w.z; + kw[j][4 * c + 3] = w.w; + } + ksc[j] = valid ? __ldg(reinterpret_cast(p.ksc + page * p.ksc_blk + n * 8)) : make_uint2(0, 0); + } + bool const valid = blk * kPage + tid < kv_len; + uint8_t vsc[8]; + if (valid) + { + uint8_t const* vb = p.vsc + page * p.vsc_blk + (tid >> 2) * 32 + (tid & 3); +#pragma unroll + for (int c = 0; c < 8; ++c) + vsc[c] = __ldg(vb + 4 * c); + } + dequant_row(s.vs + tid * kPitch, p.v + page * p.v_blk + tid * 64, vsc, valid); + } + __syncthreads(); + + if (active) + { + float const scale_log2 = p.sm_scale * 1.4426950409F * __ldg(p.kgs); + int const q4 = qd * 8; + // S^T = K q^T for this warp's 32 tokens: rows (tokens) g / g + 8 of each m-tile, columns (heads) 2 qd + {0, 1}. + float sacc[2][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + sacc[j][i] = 0.F; +#pragma unroll + for (int ks2 = 0; ks2 < kDim / 16; ks2 += 2) + { + uint32_t bq[4]; + ldsm_x4(bq, smem_u32(s.qs + (lane & 7) * kPitch + ks2 * 16 + (lane >> 3) * 8)); +#pragma unroll + for (int u = 0; u < 2; ++u) + { + int const ks = ks2 + u; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) + { + uint32_t a[4]; +#pragma unroll + for (int hi = 0; hi < 2; ++hi) + { + int const j = 2 * mt + hi; + uint32_t const scw = ks < 4 ? ksc[j].x : ksc[j].y; + uint8_t const sc = static_cast(scw >> (8 * (ks & 3))); + uint32_t const sh = e4m3x2_to_f16x2(static_cast(sc | (sc << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); + uint32_t const r0 = e2m1x2_to_f16x2((kw[j][2 * ks] >> q4) & 0xFF); + uint32_t const r1 = e2m1x2_to_f16x2((kw[j][2 * ks + 1] >> q4) & 0xFF); + __half2 const h0 = __hmul2(*reinterpret_cast<__half2 const*>(&r0), scale); + __half2 const h1 = __hmul2(*reinterpret_cast<__half2 const*>(&r1), scale); + a[hi] = *reinterpret_cast(&h0); + a[2 + hi] = *reinterpret_cast(&h1); + } + mma(sacc[mt], a, bq[2 * u], bq[2 * u + 1]); + } + } + } + float mloc[2] = {kNegInf, kNegInf}; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const t = warp * 32 + mt * 16 + g + 8 * (i >> 1); + float const v = blk * kPage + t < kv_len ? 0.F + sacc[mt][i] * scale_log2 : kNegInf; + sacc[mt][i] = v; + mloc[i & 1] = fmaxf(mloc[i & 1], v); + } +#pragma unroll + for (int e = 0; e < 2; ++e) + { + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 4)); + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 8)); + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 16)); + } + if (g == 0) + { + s.red[warp][2 * qd] = mloc[0]; + s.red[warp][2 * qd + 1] = mloc[1]; + } + __syncthreads(); + // Every thread reduces its two heads' four warp maxima itself (the same expression, so the same bits): no + // shared row-max stage and its second barrier. +#pragma unroll + for (int e = 0; e < 2; ++e) + { + int const row = 2 * qd + e; + m_row[e] = fmaxf(fmaxf(s.red[0][row], s.red[1][row]), fmaxf(s.red[2][row], s.red[3][row])); + } + // A row with no valid token has m = -inf: its probabilities are ex2(-inf - 0) = 0, never NaN. + float const m_use[2] = {m_row[0] == kNegInf ? 0.F : m_row[0], m_row[1] == kNegInf ? 0.F : m_row[1]}; + float lloc[2] = {0.F, 0.F}; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + float const pv = ex2(sacc[mt][i] - m_use[i & 1]); + lloc[i & 1] += pv; + int const t = warp * 32 + mt * 16 + g + 8 * (i >> 1); + s.ps[(2 * qd + (i & 1)) * kPitch + t] = __float2half_rn(pv); + } +#pragma unroll + for (int e = 0; e < 2; ++e) + { + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 4); + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 8); + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 16); + } + __syncthreads(); // row-max reads of red done before red is reused; ps complete + if (g == 0) + { + s.red[warp][2 * qd] = lloc[0]; + s.red[warp][2 * qd + 1] = lloc[1]; + } + // O^T[warp's 32 dims, 8 heads] = V^T P^T over all 128 tokens. +#pragma unroll + for (int kk = 0; kk < kPage; kk += 32) + { + uint32_t b[4]; + ldsm_x4(b, smem_u32(s.ps + (lane & 7) * kPitch + kk + (lane >> 3) * 8)); +#pragma unroll + for (int u = 0; u < 2; ++u) + { +#pragma unroll + for (int mt = 0; mt < 2; ++mt) + { + uint32_t a[4]; + int const t = kk + u * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + int const d = warp * 32 + mt * 16 + ((lane >> 3) & 1) * 8; + ldsm_x4_t(a, smem_u32(s.vs + t * kPitch + d)); + mma(acc[mt], a, b[2 * u], b[2 * u + 1]); + } + } + } + __syncthreads(); + // Likewise the row sums, in the same order (red is not written again: the merge uses ex / wgt). +#pragma unroll + for (int e = 0; e < 2; ++e) + { + int const row = 2 * qd + e; + l_row[e] = (s.red[0][row] + s.red[1][row]) + (s.red[2][row] + s.red[3][row]); + } + } + + // Normalize; push (head, dim) to CTA 2 * head + dim / 64, slot recv[rank][dim % 64], plus each head's LSE (log2) + // to both of its CTAs. Posted DSMEM stores, made visible by the cluster barrier. + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + { + float const vgs = __ldg(p.vgs); + float scale[2], lse[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) + { + if constexpr (ONE_DIV) + { + // Unnormalized: the merge divides once. An inactive CTA pushes m = -inf, l = 0 (weight 0). + scale[e] = 1.F; + lse[e] = m_row[e]; + } + else if constexpr (LINEAR) + { + scale[e] = l_row[e] > 0.F ? 1.0F / l_row[e] : 0.F; + lse[e] = l_row[e] > 0.F ? m_row[e] + lg2(l_row[e]) : kNegInf; + } + else + { + float const lse_i + = active && m_row[e] > kNegInf ? m_row[e] + lg2(ex2(kNegInf - m_row[e]) + l_row[e]) : kNegInf; + scale[e] = lse_i > kNegInf ? ex2(m_row[e] - lse_i) : 0.F; + lse[e] = lse_i; + } + scale[e] = scale[e] * vgs; + } + uint32_t dst[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) + dst[e] = map_cluster(&s.recv[rank][0], (2 * qd + e) * 2 + (warp >> 1)); +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dl = (warp & 1) * 32 + mt * 16 + g + 8 * (i >> 1); + st_cluster_f32(dst[i & 1] + dl * 4, acc[mt][i] * scale[i & 1]); + } + if (warp == 0 && g == 0) + { +#pragma unroll + for (int e = 0; e < 2; ++e) +#pragma unroll + for (int hf = 0; hf < 2; ++hf) + { + st_cluster_f32(map_cluster(&s.recv_lse[rank], (2 * qd + e) * 2 + hf), lse[e]); + if constexpr (ONE_DIV) + { + st_cluster_f32(map_cluster(&s.recv_l[rank], (2 * qd + e) * 2 + hf), l_row[e]); + } + } + } + } + cluster_sync(); + + // CTA `rank` merges head rank / 2, dims (rank % 2) * 64 + dl; thread = (chunk half, dl). + int const h = static_cast(rank) >> 1; + int const dh = static_cast(rank) & 1; + int const dl = tid & (kHalfDim - 1); + int const hc = tid >> 6; + if constexpr (ONE_DIV) + { + // Thread c < 16 takes chunk c's weight 2^(m_c - max m) (wgt) and its weighted row sum (ex); the output is + // the weighted partials over the weighted row sums, both summed in chunk order, one division. + if (tid < kTopk) + { + float m_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + m_max = fmaxf(m_max, s.recv_lse[c]); + m_max = m_max == kNegInf ? 0.F : m_max; + float const w = ex2(s.recv_lse[tid] - m_max); + s.wgt[tid] = w; + s.ex[tid] = w * s.recv_l[tid]; + } + __syncthreads(); + } + else + { + // The merge weights (exp2 of each chunk's LSE minus the largest, divided by their sum in chunk order) are + // formed once per CTA: thread c < 16 takes chunk c's exponential, then divides it by the sum each of those + // threads forms in the same order. Same values in the same order as every thread forming all 16: bitwise the + // same rows. + if (tid < kTopk) + { + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + lse_max = fmaxf(lse_max, s.recv_lse[c]); + lse_max = lse_max == kNegInf ? 0.F : lse_max; + s.ex[tid] = ex2(s.recv_lse[tid] - lse_max); + } + __syncthreads(); + if (tid < kTopk) + { + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + denom += s.ex[c]; + denom = denom > 0.F ? denom : 1.F; + s.wgt[tid] = __fdiv_rn(s.ex[tid], denom); + } + __syncthreads(); + } + float o = 0.F; +#pragma unroll + for (int c = 0; c < kTopk / 2; ++c) + { + int const cc = hc * (kTopk / 2) + c; + o += s.recv[cc][dl] * s.wgt[cc]; + } + if (hc) + s.part[dl] = o; + __syncthreads(); + if (!hc) + { + o += s.part[dl]; + if constexpr (ONE_DIV) + { + float den = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + den += s.ex[c]; + o = den > 0.F ? __fdiv_rn(o, den) : 0.F; + } + int const col = h * kDim + dh * kHalfDim + dl; + __nv_bfloat16 const ob = __float2bfloat16_rn(o); + p.out[tok * p.out_stride + col] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + float const xq = x * qs; + uint16_t q2; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2) : "f"(0.F), "f"(xq)); + p.q8[static_cast(tok) * (p.num_sf_k_tiles * 128) + col] = static_cast(q2 & 0xFF); + if (lane == 0) + { + int const k_sf = col / 32; + int64_t const off = static_cast(tok / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (tok % 32) * 16 + ((tok % 128) / 32) * 4 + k_sf % 4; + p.sf[off] = static_cast(e); + } + } + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +cudaError_t configure() +{ + auto kernel = m3_tp8_lc8_attn_kernel; + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes())); + if (err != cudaSuccess) + { + return err; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); +} + +template +cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) +{ + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs + // every variant eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + auto kernel = m3_tp8_lc8_attn_kernel; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(num_tokens * kTopk); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smemBytes(); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kTopk; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +template +cudaError_t launchPrewait(Params const& prm, int num_tokens, bool prewait, bool use_pdl, cudaStream_t stream) +{ + if (prm.one_div) + { + return prewait ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); + } + return prewait ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); +} + +} // namespace + +cudaError_t invokeSparseDecodeLc8m(MinimaxM3Tp8SparseDecodeLc8Params const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream) +{ + if (linearSoftmax) + { + return mxfp8 ? launchPrewait(params, numTokens, prewait, usePdl, stream) + : launchPrewait(params, numTokens, prewait, usePdl, stream); + } + return mxfp8 ? launchPrewait(params, numTokens, prewait, usePdl, stream) + : launchPrewait(params, numTokens, prewait, usePdl, stream); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h new file mode 100644 index 000000000000..7f59a0f7ff59 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h @@ -0,0 +1,105 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +//! Operands of the MiniMax-M3 TP8 one-launch long-context sparse decode. Strides are in elements of the pointed-to +//! type. Geometry is fixed: 8 query heads on one KV head, head_dim 128, page 128, top-k 16. +struct MinimaxM3Tp8SparseDecodeLc8Params +{ + uint8_t const* q; // e4m3 [T, 8, 128] + int64_t q_tok, q_head; + uint8_t const* k; // packed e2m1 [pages, 1, 128, 64] + uint8_t const* v; + int64_t k_blk, v_blk; + uint8_t const* ksc; // e4m3 [pages, 1, 128, 8], token-major + uint8_t const* vsc; // e4m3, 4x4 token-quad order + int64_t ksc_blk, vsc_blk; + float const* kgs; + float const* vgs; + int const* topk; // [T, 16] view (kv head 0) + int64_t t_tok, t_k; + int const* block_table; + int64_t bt_stride; + int const* seq_lens; + int dql; + float sm_scale; + __nv_bfloat16* out; // [T, 1024] + int64_t out_stride; + uint8_t* q8; // e4m3 [rows, num_sf_k_tiles * 128] + uint8_t* sf; + int num_sf_k_tiles; + // Optional L2 prefetch of two read-only ranges (the layer's o_proj weight and scales), split across the grid's + // CTAs and issued before the grid dependency: a hint, so the output does not depend on it. nullptr = none. + uint8_t const* pf0; + int64_t pf0_bytes; + uint8_t const* pf1; + int64_t pf1_bytes; + // Optional next-layer KV prefetch: CTA (token, r) records its selected block id in this layer's [rows, 16] id + // buffer (n1_self) and pulls into L2 the K/V pages and scales of the block the next sparse layer's CTA (token, r) + // selected at the previous step (n1_next, from that layer's pools n1_k / n1_v / n1_ksc / n1_vsc, this layer's + // strides). Hints only. nullptr = off. + int* n1_self; + int const* n1_next; + uint8_t const* n1_k; + uint8_t const* n1_v; + uint8_t const* n1_ksc; + uint8_t const* n1_vsc; + int64_t n1_pages; // pages in the next layer's pool; a recorded block outside it is skipped + // Identity selection (short-sequence steps, every request within 16 blocks): CTA r attends block r of its + // token's request when r < min(16, blocks), from the device sequence length; topk is not read. + int identity; + // The step is an identity step (short sequences: no scorer / select ran this step), whether or not this kernel + // derives its blocks from the lengths (identity) or reads the staged table: q is then read after the grid wait, + // because the predecessor is the producer, which releases its dependents at entry, before it writes q. + int identity_step; + // One-division merge (MSA_ONE_DIV_MERGE): the CTAs push unnormalized partials with their row max and sum, and the + // merge divides once per output. Deterministic, a different rounding sequence from the per-chunk normalization. + int one_div; +}; + +//! One cluster of 16 CTAs per query token: CTA r attends the token's r-th selected block with transposed MMAs (the 8 +//! heads fill the MMA n dimension), then CTA r merges (head r / 2, dims (r % 2) * 64 + [0, 64)) of the 16 partials +//! over DSMEM into the BF16 [T, 1024] row (and, when mxfp8, its MXFP8 copy with swizzled UE8M0 scales). Dependents +//! are released right after the grid dependency wait. +//! +//! With prewait, the request's seq_lens entry, its q rows (widened to FP16) and its block-table row (up to 8192 +//! entries, in dynamic shared memory; longer rows are read from global memory) are staged before the grid dependency +//! wait; outputs are bit-identical to prewait=false. +//! +//! Instantiated variants: linearSoftmax x mxfp8 x prewait, all in {false, true}. The device code is SM100-family +//! only; on other architectures the kernel compiles to an empty body, so callers check the device first. +cudaError_t invokeSparseDecodeLc8m(MinimaxM3Tp8SparseDecodeLc8Params const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu new file mode 100644 index 000000000000..fd101e7f521c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kThreads = 256; + +__global__ void __launch_bounds__(kThreads) + specOffsetsKernel(int32_t* __restrict__ pos, int32_t const* __restrict__ posOff, int64_t nPos, + int32_t* __restrict__ kv, int32_t const* __restrict__ kvOff, int64_t nKv) +{ + for (int64_t i = threadIdx.x; i < nPos; i += kThreads) + { + pos[i] = static_cast(static_cast(pos[i]) + static_cast(posOff[i])); + } + for (int64_t j = threadIdx.x; j < nKv; j += kThreads) + { + kv[j] = static_cast(static_cast(kv[j]) + static_cast(kvOff[j])); + } +} + +} // namespace + +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream) +{ + specOffsetsKernel<<<1, kThreads, 0, stream>>>(pos, posOff, nPos, kv, kvOff, nKv); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h new file mode 100644 index 000000000000..3cfd05b73369 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +//! The overlap scheduler's speculative step correction in one launch: pos[i] += posOff[i] for i < nPos and +//! kv[j] += kvOff[j] for j < nKv (int32, wrapping as the ATen adds do). A plain launch: it is a step-input writer that +//! pre-wait readers rely on (rubin-20261015/PREWAIT_AUDIT.md item b). +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu new file mode 100644 index 000000000000..b1ab46a36088 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kThreads = 512; +constexpr int kWarps = kThreads / 32; + +// One warp per copy. Plain (coherent) loads: a device copy may read bytes another warp of this CTA just stored, which +// the CTA barrier between the two kinds makes visible; a non-coherent load could return the stale line. +__device__ __forceinline__ void copyWarp(uint8_t* dst, uint8_t const* src, int64_t bytes, int lane) +{ + auto const d = reinterpret_cast(dst); + auto const s = reinterpret_cast(src); + if (((d | s | static_cast(bytes)) & 15) == 0) + { + auto* d16 = reinterpret_cast(dst); + auto const* s16 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 16; i += 32) + { + d16[i] = s16[i]; + } + } + else if (((d | s | static_cast(bytes)) & 3) == 0) + { + auto* d4 = reinterpret_cast(dst); + auto const* s4 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 4; i += 32) + { + d4[i] = s4[i]; + } + } + else + { + for (int64_t i = lane; i < bytes; i += 32) + { + dst[i] = src[i]; + } + } +} + +// One CTA: a step's window is a few dozen copies of a few hundred bytes, so the launch is the cost; a CTA barrier +// orders the host copies before the device copies where the parallel scatter needed a second launch. +__global__ void __launch_bounds__(kThreads) stepPrologueKernel(uint8_t* window, int64_t* epoch, int64_t maxDesc) +{ + auto* header = reinterpret_cast(window); + int64_t const numHost = min(header[0], maxDesc); + int64_t const numDevice = min(header[1], maxDesc); + int64_t const* desc = header + 2; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (epoch != nullptr && threadIdx.x == 0) + { + *epoch += 1; + } + for (int64_t i = warp; i < numHost; i += kWarps) + { + int64_t const* c = desc + 3 * i; + copyWarp(reinterpret_cast(c[0]), window + c[1], c[2], lane); + } + __syncthreads(); + for (int64_t i = warp; i < numDevice; i += kWarps) + { + // The device copies' slots follow all header[0] host copies, as the parallel scatter indexes them. + int64_t const* c = desc + 3 * (header[0] + i); + copyWarp(reinterpret_cast(c[0]), reinterpret_cast(c[1]), c[2], lane); + } + // Every thread has read the header and the descriptors: consume the window. + __syncthreads(); + if (threadIdx.x < 2) + { + header[threadIdx.x] = 0; + } +} + +} // namespace + +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream) +{ + stepPrologueKernel<<<1, kThreads, 0, stream>>>(window, epoch, maxDesc); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h new file mode 100644 index 000000000000..c125dede4361 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h @@ -0,0 +1,47 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +//! The decode graph's step-input prologue in one launch (step_inputs.StepInputPack with early_copy): the work of the +//! two parallel scatter launches and the header clear that follow them, in that order. +//! +//! The window is the pack's staging layout (kernels/stepInputsScatter.h): an int64 header `[numHostCopies, +//! numDeviceCopies]`, then one `(dst, src, bytes)` int64 triple per copy, host copies first; a host copy's `src` is a +//! byte offset into the window, a device copy's a device address. Up to `maxDesc` copies of each kind are written (the +//! slots the parallel scatter's grid covers), every host copy before any device copy, so a device copy may read a +//! destination a host copy filled; the copies of one kind must not overlap (the pack issues the window only then). +//! Afterwards the header is zeroed, so a replay with no new upload writes nothing. With `epoch` non-null, it is +//! incremented once. +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu new file mode 100644 index 000000000000..4fad864f2d64 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kThreads = 256; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves): the +// instruction sequence of the Triton kernel at limit 7, alpha 1.702, beta 1 (Indexer kdev17, bitwise on SM107): the +// gate clamp as min.bf16x2 before the widening, the up clamp as min.xorsign.abs, (u + 1), (g * -1.702) * log2(e), +// ex2.approx, + 1, div.full 1 / that, s * g, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t swigluOaiBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, %1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +template +__global__ void __launch_bounds__(kThreads) + swigluOaiKernel(uint32_t const* __restrict__ x, uint32_t* __restrict__ out, int64_t rows, int64_t halfD) +{ + if constexpr (kPdl) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int64_t const i = blockIdx.x * static_cast(kThreads) + threadIdx.x; + if (i >= rows * halfD) + { + return; + } + int64_t const row = i / halfD; + int64_t const col = i - row * halfD; + uint32_t const* xr = x + row * 2 * halfD; + out[i] = swigluOaiBf16x2(xr[col], xr[halfD + col]); +} + +} // namespace + +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream) +{ + int64_t const n = rows * halfD; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast((n + kThreads - 1) / kThreads)); + cfg.blockDim = dim3(kThreads); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return pdl ? cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD) + : cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h new file mode 100644 index 000000000000..78033e018997 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +//! The dense MLP's swigluoai activation: out[m, j] = bf16(min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + +//! 1)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both contiguous, d even): bit for bit +//! the Triton silu_and_mul kernel at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With `pdl` the launch allows +//! programmatic serialization and the grid waits on its dependency, then releases its dependents before it computes +//! (the down GEMV stages its weights during the activation and reads it only after its own grid wait). +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/quantization.cu b/cpp/tensorrt_llm/kernels/quantization.cu index f6cf68053094..33a3801a5232 100644 --- a/cpp/tensorrt_llm/kernels/quantization.cu +++ b/cpp/tensorrt_llm/kernels/quantization.cu @@ -194,6 +194,9 @@ void invokeFP4Quantization(int b, int m, int n, T const* input, float const* SFS //////////////////////////////////////////////////////////////////////////////////////////////////// // MXFP8 Quantization +// Kill switch for the SM107 quantize_mxfp8_rows path of invokeMxFP8Quantization. +constexpr bool kMxfp8QuantSm107Rows = true; + template void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream) @@ -207,6 +210,33 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, return; } + // SM107: one unpadded matrix with swizzled K32 scales runs quantize_mxfp8_rows (same math, one column chunk per + // thread, 8 rows in flight); bitwise the same output. + if constexpr (SF_VEC_SIZE == 32 && SF_OUTPUT_VEC_SIZE == 32) + { + int const colThreads = n / CVT_ELTS_PER_THREAD; + if (kMxfp8QuantSm107Rows && tensorrt_llm::common::getSMVersion() == 107 && b == 1 && n == padded_n + && layout == QuantizationSFLayout::SWIZZLED && n % 128 == 0 && colThreads <= 1024 && colThreads % 32 == 0) + { + constexpr int kRows = 8; + int const ctasPerSm = colThreads > 512 ? 1 : 2; + dim3 grid(std::min(int(PadUpFn(m, 128) + kRows - 1) / kRows, multiProcessorCount * ctasPerSm)); + cudaLaunchConfig_t config; + config.gridDim = grid; + config.blockDim = dim3(colThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, quantize_mxfp8_rows, m, n, input, + reinterpret_cast(output), reinterpret_cast(SFOuput))); + return; + } + } + // Grid, Block size. // Each thread converts 8 values. dim3 block(std::min(int(padded_n / CVT_ELTS_PER_THREAD), 512)); diff --git a/cpp/tensorrt_llm/kernels/quantization.cuh b/cpp/tensorrt_llm/kernels/quantization.cuh index 4b986cfe3b82..0db2258d8a46 100644 --- a/cpp/tensorrt_llm/kernels/quantization.cuh +++ b/cpp/tensorrt_llm/kernels/quantization.cuh @@ -974,6 +974,64 @@ quantize_with_block_size( #endif } +//! MXFP8 (K32 UE8M0 scales, swizzled 128x4 layout) quantization of one unpadded [numRows, numCols] matrix with the +//! per-element math of quantize_with_block_size (cvt_warp_fp16_to_mxfp8): one 8-element column chunk per thread +//! (blockDim.x == numCols / 8) and ROWS rows loaded before any is converted, so more bytes are in flight per thread. +//! Requires numCols % 128 == 0 and numCols / 8 <= 1024 a multiple of 32. +template +__global__ void __launch_bounds__(1024) + quantize_mxfp8_rows(int32_t numRows, int32_t numCols, Type const* in, uint64_t* out, uint32_t* SFout) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using PackedVec = PackedVec; + static constexpr int CVT_NUM_THREADS_PER_SF = 32 / CVT_ELTS_PER_THREAD; + int const colIdx = threadIdx.x; + int const numColThreads = numCols / CVT_ELTS_PER_THREAD; + int const numPaddedRowsForSf = PadUpFn(numRows, 128); + std::optional const optionalBatchIdx = 0; + std::optional const optionalNumRows = numRows; + + cudaGridDependencySynchronize(); + for (int rowBase = blockIdx.x * ROWS; rowBase < numPaddedRowsForSf; rowBase += gridDim.x * ROWS) + { + PackedVec inVecs[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; r++) + { + if (rowBase + r < numRows) + { + inVecs[r] = reinterpret_cast( + in)[static_cast(rowBase + r) * numColThreads + colIdx]; + } + } +#pragma unroll + for (int r = 0; r < ROWS; r++) + { + int const rowIdx = rowBase + r; + if (rowIdx >= numPaddedRowsForSf) + { + break; + } + auto sf_out = cvt_quant_get_sf_out_offset( + optionalBatchIdx, rowIdx, colIdx, optionalNumRows, numCols / 32, SFout, QuantizationSFLayout::SWIZZLED); + if (rowIdx < numRows) + { + out[static_cast(rowIdx) * numColThreads + colIdx] + = cvt_warp_fp16_to_mxfp8(inVecs[r], sf_out); + } + else + { + cvt_store_replicated_sf<1>(sf_out, 0x00); + } + } + } + // See quantize_with_block_size: every thread's stores are visible before any thread triggers completion. + __threadfence(); + __syncthreads(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + __global__ void block_scale_interleave_kernel( int numbatches, int numRows, int numCols, uint8_t const* SFIn, uint8_t* SFOutput); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 79cc1eb5dd7e..8af113270f1e 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -254,7 +254,7 @@ class TllmGenFmhaKernel } } - static bool shouldUseNvrtc(FmhaOptions const& options) + bool shouldUseNvrtc(FmhaOptions const& options) const { #ifdef TLLM_RUBIN_FEATURES // Spcompress kernels must use precompiled cubins. @@ -263,7 +263,20 @@ class TllmGenFmhaKernel return false; } #endif // TLLM_RUBIN_FEATURES - // Sparse MQA/GQA uses NVRTC path for now because no model really uses it. + // The cubin table has no dense paged-KV kernels for 128-token pages, and the page size is + // part of the kernel hash. On SM107, generate those BF16/E4M3 context and KeepsMmaAb + // generation kernels for the actual cache layout instead. + constexpr int kPageSize128 = 128; + if (mSM == kSM_107 + && (isContextKernel(options.mFmhaKernelType) || isKeepsMmaAbForGenerationKernel(options.mFmhaKernelType)) + && !options.mIsMlaGen && !isTokenSparse(options.mSparseType) && isPagedKv(options.mQkvLayout) + && options.mNumTokensPerPage == kPageSize128 && options.mDtypeQ == options.mDtypeKv + && (options.mDtypeQ == tg::Dtype::Bfloat16 || options.mDtypeQ == tg::Dtype::E4m3) + && options.mDtypeOut == tg::Dtype::Bfloat16) + { + return true; + } + // Sparse MQA/GQA uses NVRTC path for now because no model really uses it. if (isStaticTokenSparse(options.mSparseType) && !options.mIsMlaGen) { return true; @@ -375,6 +388,20 @@ class TllmGenFmhaKernel 1152, 1280, 1408, 1536, 1664, 1792, 1920, 2048, 2176, 2304, 2432, 2560, 2688, 2816, 2944, 3072, 3328, 3584, 3840, 4096, 4352, 4608, 4864, 5120, 5376, 5632, 5888, 6144, 6656, 7168, 7680, 8192, 8704, 9216, 9728, 10240, 10752, 11264, 11776, 12288, 13312, 14336, 16384, 20480, 24576, 32768, 40960, 49152, 65536, 98304, 131072}; + // The same grid continued past 131072 in 32768-token steps to 1M tokens: kernel selection still changes there + // (e.g. MultiCtasKv vs MultiCtasKvCga), so a long-context step whose + // KV length falls between two grid points otherwise compiles its kernel inside a served request (~4-5 s). + // Selection is a cheap heuristic and compiled kernels are cached by config, so the finer grid only adds the + // compiles of configs the coarser grid missed. Kill: TRTLLM_FMHA_LONG_WARMUP=0. + inline static std::vector const kLongWarmupSeqLenKvCandidates = [] + { + std::vector sizes = kDefaultWarmupSeqLenKvCandidates; + for (int size = 131072 + 32768; size <= 1048576; size += 32768) + { + sizes.push_back(size); + } + return sizes; + }(); // Selection of seqLenQ: typically our prefill kernels would not use JIT and seqLenQ is fixed during decode. // However seqLenQ can change when we use generation kernel to do prefill. SeqLenQ is sensitive in the range 1-16. // Below sequence are chosen to cover most sparse attention cases. @@ -411,6 +438,11 @@ class TllmGenFmhaKernel std::tie(options, optionsFromArgs, ctaDim) = selectKernelAutotuned(options, optionsFromArgs, params.mMultiProcessorCount); + // Same override as run(), so the warmed config matches the dispatched one. + if (mNumEltsPerSageAttnBlkP + mNumEltsPerSageAttnBlkV > 0) + { + options.mTileScheduler = TileScheduler::Static; + } checkFmhaOptions(options, optionsFromArgs); updateFmhaOptions(options, optionsFromArgs); @@ -421,10 +453,12 @@ class TllmGenFmhaKernel if (shouldUseNvrtc(options)) { + // Compile exactly what run() compiles (same target arch), so the warmed kernel is the one run() hits. + FmhaOptions const nvrtcOptions = nvrtcCompileOptions(options); FmhaConfig fmhaConfig; - fmhaConfig.mOptions = options; + fmhaConfig.mOptions = nvrtcOptions; std::ostringstream sstream; - populateJsonConfig(options, sstream); + populateJsonConfig(nvrtcOptions, sstream); fmhaConfig.mGenCfgJsonStr = sstream.str(); fmhaConfig.mExecPath = getExecPath().c_str(); @@ -442,9 +476,67 @@ class TllmGenFmhaKernel } } + // Context-phase grid: the kernel variant (e.g. static vs persistent tile scheduler, tile shape) follows the CTA + // count, i.e. batch x ceil(seqLenQ / tileQ), so cover each power of two and the tile boundary just past it + // (chunked prefill runs chunk + 1 query tokens), with every KV length >= seqLenQ (cached prefix) from the KV grid. + inline static std::vector const kContextWarmupBatchSizeCandidates = {1, 2, 4, 8, 16, 32, 64, 128, 256}; + inline static std::vector const kContextWarmupSeqLenQCandidates = {1, 128, 129, 256, 257, 512, 513, 1024, 1025, + 2048, 2049, 4096, 4097, 8192, 8193, 16384, 16385, 32768, 32769, 65536, 65537, 131072}; + + void runContextJITWarmupGrid(RunnerParams const& runnerParams) + { + static bool const contextWarmup + = tensorrt_llm::common::getIntEnv("TRTLLM_FMHA_CONTEXT_WARMUP").value_or(1) != 0; + int const maxBatchSize = runnerParams.mJITWarmupMaxNumRequests; + int const maxSeqLenQ = runnerParams.mJITWarmupMaxSeqLenQ; + int const maxSeqLenKv = runnerParams.mJITWarmupMaxSeqLenKv; + if (!contextWarmup || maxBatchSize <= 0 || maxSeqLenQ <= 0 || maxSeqLenKv <= 0) + { + return; + } + static bool const longWarmup = tensorrt_llm::common::getIntEnv("TRTLLM_FMHA_LONG_WARMUP").value_or(1) != 0; + std::vector const seqLenKvCandidates = makeWarmupCandidateSizes( + longWarmup ? kLongWarmupSeqLenKvCandidates : kDefaultWarmupSeqLenKvCandidates, maxSeqLenKv); + auto warmupParams = runnerParams; + for (int batchSize : makeWarmupCandidateSizes(kContextWarmupBatchSizeCandidates, maxBatchSize)) + { + warmupParams.mBatchSize = batchSize; + for (int seqLenQ : makeWarmupCandidateSizes(kContextWarmupSeqLenQCandidates, maxSeqLenQ)) + { + warmupParams.mMaxSeqLenQ = seqLenQ; + for (int seqLenKv : seqLenKvCandidates) + { + if (seqLenKv < seqLenQ) + { + continue; + } + warmupParams.mMaxSeqLenKv = seqLenKv; + warmupParams.mSumOfSeqLensQ = static_cast( + std::min(static_cast(batchSize) * seqLenQ, std::numeric_limits::max())); + warmupParams.mSumOfSeqLensKv = static_cast( + std::min(static_cast(batchSize) * seqLenKv, std::numeric_limits::max())); + warmupOneKernel(warmupParams); + } + } + } + } + void runJITWarmupGridIfRequested(RunnerParams const& runnerParams) { - if (!runnerParams.mJITWarmup || runnerParams.mKernelType != FmhaKernelType::Generation) + if (!runnerParams.mJITWarmup) + { + return; + } + if (isContextKernel(runnerParams.mKernelType)) + { + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + TLLM_CUDA_CHECK(cudaStreamIsCapturing(runnerParams.stream, &captureStatus)); + TLLM_CHECK_WITH_INFO(captureStatus == cudaStreamCaptureStatusNone, + "TRTLLM-Gen FMHA JIT warmup must not run during CUDA graph capture."); + runContextJITWarmupGrid(runnerParams); + return; + } + if (runnerParams.mKernelType != FmhaKernelType::Generation) { return; } @@ -471,7 +563,9 @@ class TllmGenFmhaKernel std::vector seqLenQCandidates = useGenKernelForPrefill ? makeWarmupCandidateSizes(kDefaultWarmupSeqLenQCandidates, maxSeqLenQ) : std::vector{runnerParams.mMaxSeqLenQ}; - std::vector seqLenKvCandidates = makeWarmupCandidateSizes(kDefaultWarmupSeqLenKvCandidates, maxSeqLenKv); + static bool const longWarmup = tensorrt_llm::common::getIntEnv("TRTLLM_FMHA_LONG_WARMUP").value_or(1) != 0; + std::vector seqLenKvCandidates = makeWarmupCandidateSizes( + longWarmup ? kLongWarmupSeqLenKvCandidates : kDefaultWarmupSeqLenKvCandidates, maxSeqLenKv); auto warmupParams = runnerParams; @@ -581,14 +675,7 @@ class TllmGenFmhaKernel if (shouldUseNvrtc(options)) { // nvrtc path - uses mFmhaInterface member for kernel caching - FmhaOptions nvrtcOptions = options; -#ifdef TLLM_RUBIN_FEATURES - // On Rubin, compile non-spcompress NVRTC kernels with Sm100f target for compatibility. - if (tg::isArchRubin(options.mCudaArch) && !options.mUsesSpcompress) - { - nvrtcOptions.mCudaArch = tg::CudaArch::Sm100f; - } -#endif // TLLM_RUBIN_FEATURES + FmhaOptions const nvrtcOptions = nvrtcCompileOptions(options); FmhaConfig fmhaConfig; fmhaConfig.mOptions = nvrtcOptions; @@ -664,6 +751,20 @@ class TllmGenFmhaKernel } private: + // The options NVRTC compiles with; shared by run() and warmupOneKernel() so both produce the same kernel. + FmhaOptions nvrtcCompileOptions(FmhaOptions const& options) const + { + FmhaOptions nvrtcOptions = options; +#ifdef TLLM_RUBIN_FEATURES + // On Rubin, compile non-spcompress NVRTC kernels with Sm100f target for compatibility. + if (tg::isArchRubin(options.mCudaArch) && !options.mUsesSpcompress) + { + nvrtcOptions.mCudaArch = tg::CudaArch::Sm100f; + } +#endif // TLLM_RUBIN_FEATURES + return nvrtcOptions; + } + // Autotuner entry for checkIfKernelExist, warmupOneKernel, and run. SM107 may select // CgaSmemReduction configs where clusterDimX * mMaxNumCtasPerSeqKv must stay <= 16; // the pre/post clamps below enforce that. Other architectures keep the legacy path. diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp index 059d29e43d85..6e7a6f1fbef9 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp @@ -43,6 +43,14 @@ void checkSmVersion(char const* op) TORCH_CHECK(sm == 100 || sm == 103, op, " requires SM100 or SM103, got SM", sm); } +// The TF32 tensor-core router also runs on SM107 (Rubin; its kernels are sm_100f-family cubins), there only with +// 16-CTA clusters: one reduction order on every rank (minimaxM3RouterTcCluster16 lets callers fall back otherwise). +void checkSmVersionTc(char const* op) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 100 || sm == 103 || sm == 107, op, " requires SM100, SM103 or SM107, got SM", sm); +} + void checkActivation(torch::Tensor const& x, torch::Tensor const& w, int64_t maxM) { TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.stride(1) == 1, @@ -96,7 +104,7 @@ std::vector minimaxM3RouterGemvFp4( std::vector minimaxM3RouterTc( torch::Tensor const& x, torch::Tensor const& w, std::optional const& sfScale, bool pdl) { - checkSmVersion("minimax_m3_router_tc"); + checkSmVersionTc("minimax_m3_router_tc"); TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.stride(1) == 1, "x must be a 2-D bf16 CUDA tensor with unit column stride"); int64_t const M = x.size(0), K = x.size(1); @@ -105,6 +113,10 @@ std::vector minimaxM3RouterTc( "minimax_m3_router_tc takes 1..48 rows and a contiguous fp32 [128, K] weight"); TORCH_CHECK(K % 1024 == 0 && (x.stride(0) * 2) % 16 == 0, "K must be a multiple of 1024 with 16-byte rows"); c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(tensorrt_llm::common::getSMVersion() != 107 + || tensorrt_llm::kernels::minimaxM3RouterTcCluster(static_cast(K)) == 16, + "minimax_m3_router_tc on SM107 runs only with 16-CTA clusters; this device cannot co-schedule them (use the " + "stock router)"); auto out = torch::empty({M, w.size(0)}, x.options().dtype(at::kFloat)); torch::Tensor q, sf; uint8_t* qp = nullptr; @@ -130,6 +142,14 @@ std::vector minimaxM3RouterTc( return {out}; } +//! Whether `minimax_m3_router_tc` runs on this device with 16-CTA clusters for hidden size `k` (device-static; the +//! SM107 caller falls back to the stock router otherwise). +bool minimaxM3RouterTcCluster16(int64_t k) +{ + checkSmVersionTc("minimax_m3_router_tc_cluster16"); + return tensorrt_llm::kernels::minimaxM3RouterTcCluster(static_cast(k)) == 16; +} + } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -139,6 +159,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def("minimax_m3_router_gemv(Tensor x, Tensor w, int rows, bool pdl) -> Tensor"); m.def("minimax_m3_router_gemv_fp4(Tensor x, Tensor w, Tensor sf_scale, bool pdl) -> Tensor[]"); m.def("minimax_m3_router_tc(Tensor x, Tensor w, Tensor? sf_scale, bool pdl) -> Tensor[]"); + m.def("minimax_m3_router_tc_cluster16(int k) -> bool", &tensorrt_llm::torch_ext::minimaxM3RouterTcCluster16); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp new file mode 100644 index 000000000000..99c182799cab --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp @@ -0,0 +1,125 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace dqk = tensorrt_llm::kernels::minimaxM3Sm107Tp2::denseQkNormRopeFp8; + +//! QK RMSNorm + partial RoPE of a packed BF16 [tokens, (q + k + v) * 128] row of the MiniMax-M3 TP2 dense layers, +//! returned as a new FP8 E4M3 tensor of the same shape (V copy-cast; Q/K rounded to BF16 before E4M3, as stock). +torch::Tensor minimaxM3Sm107Tp2DenseQkNormRopeFp8(torch::Tensor const& qkv, int64_t numHeadsQ, int64_t numHeadsK, + int64_t numHeadsV, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, + double base, bool isNeox, torch::Tensor const& positionIds, bool useGemma) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); + TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); + TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); + TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + at::cuda::CUDAGuard guard(qkv.device()); + auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); + dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(rotaryDim), static_cast(eps), + qWeight.data_ptr(), kWeight.data_ptr(), static_cast(base), !isNeox, positionIds.data_ptr(), + useGemma, nullptr, at::cuda::getCurrentCUDAStream(qkv.get_device())); + return out; +} + +//! As minimaxM3Sm107Tp2DenseQkNormRopeFp8, and the first numLiveTokens rows' K and V are also stored into the layer's +//! FP8 pool kvBuffers (the HND view [pages, 2, kv heads, tokens per block, 128]) at their int32 cache slots; with +//! qCompact, every row's Q is also stored there compact ([tokens, q heads, 128]). +torch::Tensor minimaxM3Sm107Tp2DenseQkNormRopeFp8Kv(torch::Tensor const& qkv, int64_t numHeadsQ, int64_t numHeadsK, + int64_t numHeadsV, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, + double base, bool isNeox, torch::Tensor const& positionIds, bool useGemma, torch::Tensor const& kvBuffers, + torch::Tensor const& outCacheLoc, int64_t numLiveTokens, std::optional const& qCompact) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); + TORCH_CHECK(numHeadsK == numHeadsV, "the pool keeps one head count for K and V"); + TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); + TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); + TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + TORCH_CHECK(kvBuffers.is_cuda() && kvBuffers.scalar_type() == at::kFloat8_e4m3fn && kvBuffers.dim() == 5); + TORCH_CHECK(kvBuffers.size(1) == 2 && kvBuffers.size(2) == numHeadsK && kvBuffers.size(4) == dqk::kHeadDim + && kvBuffers.stride(4) == 1, + "kv_buffers must be the HND pool view [pages, 2, kv heads, tokens per block, 128]"); + TORCH_CHECK(outCacheLoc.is_cuda() && outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous()); + TORCH_CHECK(numLiveTokens >= 0 && numLiveTokens <= qkv.size(0) && numLiveTokens <= outCacheLoc.numel()); + at::cuda::CUDAGuard guard(qkv.device()); + auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); + dqk::KvCacheView view; + view.base = static_cast<__nv_fp8_e4m3*>(kvBuffers.data_ptr()); + view.outCacheLoc = outCacheLoc.data_ptr(); + view.pageStride = kvBuffers.stride(0); + view.kvStride = kvBuffers.stride(1); + view.headStride = kvBuffers.stride(2); + view.tokenStride = kvBuffers.stride(3); + view.tokensPerBlock = static_cast(kvBuffers.size(3)); + view.numLiveTokens = static_cast(numLiveTokens); + if (qCompact.has_value()) + { + TORCH_CHECK(qCompact->is_cuda() && qCompact->is_contiguous() && qCompact->scalar_type() == at::kFloat8_e4m3fn + && qCompact->numel() == qkv.size(0) * numHeadsQ * dqk::kHeadDim, + "q_compact must be a contiguous FP8 [tokens, q heads, 128] tensor"); + view.qCompact = static_cast<__nv_fp8_e4m3*>(qCompact->data_ptr()); + } + dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(rotaryDim), static_cast(eps), + qWeight.data_ptr(), kWeight.data_ptr(), static_cast(base), !isNeox, positionIds.data_ptr(), + useGemma, &view, at::cuda::getCurrentCUDAStream(qkv.get_device())); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, " + "bool use_gemma) -> Tensor"); + m.def( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, " + "bool use_gemma, Tensor(a!) kv_buffers, Tensor out_cache_loc, int num_live_tokens, Tensor(b!)? q_compact=None) " + "-> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2DenseQkNormRopeFp8); + m.impl("minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2DenseQkNormRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DraftStepOp.cpp new file mode 100644 index 000000000000..cc452ef4c776 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DraftStepOp.cpp @@ -0,0 +1,220 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h" + +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace ds = tensorrt_llm::kernels::minimaxM3Sm107Tp2::draftStep; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_sm107tp2_draft_step expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107tp2_draft_step requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +void checkInt(torch::Tensor const& t, int64_t n, char const* name) +{ + TORCH_CHECK(t.scalar_type() == at::kInt && t.is_contiguous() && t.numel() >= n, "draft_step: ", name, + " must be contiguous int32 with at least ", n, " elements"); +} + +} // namespace + +//! One fused greedy draft pass over `rows` <= 8 requests: the BF16 head (this rank's fragment-ordered shard wf +//! [N / 16, K / 16, 256], or with excOff / excPos / excW / expBase its 12-bit coding, wf uint8 [N * K * 3 / 2]; +//! global ids vocabOffset..) over x[xRows[b]] (or x[b]), the global argmax across the TP ranks through the multicast +//! slots wsMcPtr / wsUc, then draftOut[b, step] = idsNext[b] = token, posNext[b] = posSrc[xRows[b] or b] + 1 and the +//! KV-length bookkeeping (see draftStep.h). +void minimaxM3Sm107Tp2DraftStep(torch::Tensor const& x, OptTensor const& xRows, int64_t rows, torch::Tensor const& wf, + int64_t vocabOffset, torch::Tensor& keys, torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, + int64_t rank, int64_t world, torch::Tensor& draftOut, int64_t step, torch::Tensor& idsNext, + torch::Tensor const& posSrc, torch::Tensor& posNext, torch::Tensor& kvLens, OptTensor const& numAcc, + OptTensor const& seqLens, int64_t draftLen, int64_t numCtas, bool pdl, OptTensor const& outKeys, + OptTensor const& excOff, OptTensor const& excPos, OptTensor const& excW, int64_t expBase, int64_t tileLo, + int64_t tileHi, int64_t totalCtas, OptTensor const& bnd, OptTensor const& bndCnt) +{ + checkArch(x); + for (torch::Tensor const* t : std::initializer_list{ + &wf, &keys, &counters, &wsUc, &draftOut, &idsNext, &posSrc, &posNext, &kvLens}) + { + TORCH_CHECK(t->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + for (OptTensor const* t : std::initializer_list{ + &xRows, &numAcc, &seqLens, &outKeys, &excOff, &excPos, &excW, &bnd, &bndCnt}) + { + TORCH_CHECK(!t->has_value() || (*t)->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "draft_step: x must be contiguous BF16 [rows, K]"); + int const K = static_cast(x.size(1)); + TORCH_CHECK(K % 16 == 0, "draft_step: K must be a multiple of 16"); + TORCH_CHECK(rows >= 1 && rows <= ds::kMaxTokens, "draft_step: 1..", ds::kMaxTokens, " rows"); + bool const coded = excOff.has_value(); + TORCH_CHECK(coded == excPos.has_value() && coded == excW.has_value(), + "draft_step: exc_off, exc_pos and exc_w come together"); + int nTiles = 0; + if (coded) + { + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "draft_step: a coded head needs K a multiple of 64"); + TORCH_CHECK(excOff->scalar_type() == at::kInt && excOff->is_contiguous() && excOff->numel() >= 2, + "draft_step: exc_off must be int32 [tiles + 1]"); + nTiles = static_cast(excOff->numel() - 1); + TORCH_CHECK(wf.scalar_type() == at::kByte && wf.is_contiguous() + && wf.numel() == static_cast(nTiles) * 16 * K * 3 / 2, + "draft_step: a coded wf must be the contiguous uint8 stream of [tiles * 16, K]"); + TORCH_CHECK(excPos->scalar_type() == at::kInt && excPos->is_contiguous(), "draft_step: exc_pos must be int32"); + TORCH_CHECK(excW->scalar_type() == at::kFloat && excW->is_contiguous() && excW->numel() == 2 * excPos->numel(), + "draft_step: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "draft_step: exponent base out of range"); + } + else + { + TORCH_CHECK(wf.scalar_type() == at::kBFloat16 && wf.dim() == 3 && wf.is_contiguous() && wf.size(1) == K / 16 + && wf.size(2) == 256, + "draft_step: wf must be the fragment-ordered BF16 weight [N / 16, K / 16, 256]"); + nTiles = static_cast(wf.size(0)); + } + int64_t const rangeHi = tileHi > 0 ? tileHi : nTiles; + TORCH_CHECK(tileLo >= 0 && tileLo < rangeHi && rangeHi <= nTiles, "draft_step: tile range"); + TORCH_CHECK(nTiles >= 1 && numCtas >= 1 && (rangeHi - tileLo + numCtas - 1) / numCtas <= ds::kMaxTilesPerCta, + "draft_step: too many weight tiles per CTA"); + TORCH_CHECK(totalCtas == 0 || totalCtas >= numCtas, "draft_step: total_ctas"); + TORCH_CHECK(bnd.has_value() == bndCnt.has_value(), "draft_step: bnd and bnd_cnt come together"); + if (bnd.has_value()) + { + TORCH_CHECK(bnd->scalar_type() == at::kFloat && bnd->is_contiguous() && bnd->numel() >= int64_t{nTiles} * 128, + "draft_step: bnd must be zero-filled fp32 [tiles, 32, 4]"); + TORCH_CHECK(bndCnt->scalar_type() == at::kInt && bndCnt->is_contiguous() && bndCnt->numel() >= nTiles, + "draft_step: bnd_cnt must be zero-filled int32 [tiles]"); + // At least one tile per CTA: a tile is then shared by at most two CTAs (two addends in bnd). + TORCH_CHECK(rangeHi - tileLo >= numCtas, "draft_step: an equal-chunk split needs a tile per CTA"); + } + TORCH_CHECK(vocabOffset >= 0 && vocabOffset + 16LL * nTiles <= 0xFFFFFFFFLL, "draft_step: vocab offset"); + if (xRows.has_value()) + { + TORCH_CHECK(xRows->scalar_type() == at::kLong && xRows->is_contiguous() && xRows->numel() >= rows, + "draft_step: x_rows must be contiguous int64 [rows]"); + } + else + { + TORCH_CHECK(x.size(0) >= rows, "draft_step: x must hold a row per request"); + } + TORCH_CHECK(world >= 1 && rank >= 0 && rank < world, "draft_step: rank / world"); + TORCH_CHECK(keys.scalar_type() == at::kLong && keys.numel() >= rows, "draft_step: keys must be int64 [>= rows]"); + TORCH_CHECK(counters.scalar_type() == at::kInt && counters.numel() >= 2, "draft_step: counters must be int32 [2]"); + TORCH_CHECK(wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, + "draft_step: ws_uc must hold [2][8][world] slots"); + TORCH_CHECK(draftOut.scalar_type() == at::kInt && draftOut.dim() == 2 && draftOut.stride(1) == 1 + && draftOut.size(0) >= rows && step >= 0 && step < draftOut.size(1), + "draft_step: draft_out must be int32 [rows, > step] with unit column stride"); + checkInt(idsNext, rows, "ids_next"); + checkInt(posNext, rows, "pos_next"); + checkInt(kvLens, rows, "kv_lens"); + TORCH_CHECK(posSrc.scalar_type() == at::kInt && posSrc.is_contiguous(), "draft_step: pos_src must be int32"); + TORCH_CHECK(numAcc.has_value() == seqLens.has_value(), "draft_step: num_acc and seq_lens come together"); + TORCH_CHECK(!numAcc.has_value() || step == 0, "draft_step: the KV rollback (num_acc) is pass 0's"); + if (numAcc.has_value()) + { + checkInt(*numAcc, rows, "num_acc"); + checkInt(*seqLens, rows, "seq_lens"); + } + if (outKeys.has_value()) + { + TORCH_CHECK(outKeys->scalar_type() == at::kLong && outKeys->is_contiguous() && outKeys->numel() >= rows, + "draft_step: out_keys must be contiguous int64 [rows]"); + } + ds::Params p{}; + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_rows = xRows.has_value() ? reinterpret_cast(xRows->data_ptr()) : nullptr; + p.M = static_cast(rows); + p.wf = reinterpret_cast(wf.data_ptr()); + p.n_tiles = nTiles; + p.vocab_offset = static_cast(vocabOffset); + p.keys = reinterpret_cast(keys.data_ptr()); + p.counters = reinterpret_cast(counters.data_ptr()); + p.ws_uc = reinterpret_cast(wsUc.data_ptr()); + p.ws_mc = reinterpret_cast(wsMcPtr); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.out_keys = outKeys.has_value() ? reinterpret_cast(outKeys->data_ptr()) : nullptr; + p.st = ds::StepArgs{draftOut.data_ptr(), static_cast(draftOut.stride(0)), static_cast(step), + idsNext.data_ptr(), posSrc.data_ptr(), posNext.data_ptr(), kvLens.data_ptr(), + numAcc.has_value() ? numAcc->data_ptr() : nullptr, + seqLens.has_value() ? seqLens->data_ptr() : nullptr, static_cast(draftLen)}; + p.cd = ds::CodedArgs{coded ? excOff->data_ptr() : nullptr, + coded ? reinterpret_cast(excPos->data_ptr()) : nullptr, + coded ? reinterpret_cast(excW->data_ptr()) : nullptr, + static_cast((expBase << 7) | (expBase << 23))}; + p.tile_lo = static_cast(tileLo); + p.tile_hi = static_cast(tileHi); + p.total_ctas = static_cast(totalCtas); + p.bnd = bnd.has_value() ? reinterpret_cast(bnd->data_ptr()) : nullptr; + p.bnd_cnt = bnd.has_value() ? reinterpret_cast(bndCnt->data_ptr()) : nullptr; + check(ds::invokeDraftStep(p, K, static_cast(numCtas), pdl, at::cuda::getCurrentCUDAStream(x.get_device()))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// Mutations through the raw multicast address ws_mc_ptr are not visible to the schema: that buffer is owned by the +// Python workspace and never aliased by the op's tensor arguments. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp2_draft_step(Tensor x, Tensor? x_rows, int rows, Tensor wf, int vocab_offset, Tensor(a!) " + "keys, " + "Tensor(b!) counters, Tensor(c!) ws_uc, int ws_mc_ptr, int rank, int world, Tensor(d!) draft_out, int step, " + "Tensor(e!) ids_next, Tensor pos_src, Tensor(f!) pos_next, Tensor(g!) kv_lens, Tensor? num_acc, " + "Tensor(h!)? seq_lens, int draft_len, int num_ctas, bool pdl, Tensor(i!)? out_keys, Tensor? exc_off=None, " + "Tensor? exc_pos=None, Tensor? exc_w=None, int exp_base=0, int tile_lo=0, int tile_hi=0, int total_ctas=0, " + "Tensor(j!)? bnd=None, Tensor(k!)? bnd_cnt=None) " + "-> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp2_draft_step", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2DraftStep); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp new file mode 100644 index 000000000000..6b154d20d3f2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp2::drafterRopeFp8Kv; + +//! The MiniMax-M3 TP2 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h): RoPE on +//! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages +//! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. +void minimaxM3Sm107Tp2DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, + torch::Tensor const& seqLens, std::optional const& specPositionOffsets, torch::Tensor& kvPages, + torch::Tensor const& blockTables, torch::Tensor& qOut, int64_t qLen, int64_t numHeadsQ, int64_t numHeadsKv, + int64_t rotaryDim, int64_t tokensPerBlock, int64_t ropeMode, bool pdl, bool maxSharedCarveout) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + 2 * numHeadsKv) * drk::kHeadDim, "qkv is not [tokens, (Hq + 2 Hkv) * 128]"); + TORCH_CHECK(qLen >= 1 && qkv.size(0) % qLen == 0, "the rows are not whole passes of q_len rows"); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= drk::kHeadDim && rotaryDim % 16 == 0, "unsupported rotary_dim"); + TORCH_CHECK(cosSin.is_cuda() && cosSin.is_contiguous() && cosSin.scalar_type() == at::kFloat + && cosSin.numel() % rotaryDim == 0, + "cos_sin must be the fp32 [positions, rotary_dim / 2] (cos, sin) table"); + int64_t const requests = qkv.size(0) / qLen; + TORCH_CHECK(seqLens.is_cuda() && seqLens.scalar_type() == at::kInt && seqLens.numel() >= requests); + if (specPositionOffsets.has_value()) + { + TORCH_CHECK(specPositionOffsets->is_cuda() && specPositionOffsets->is_contiguous() + && specPositionOffsets->scalar_type() == at::kInt && specPositionOffsets->numel() >= requests * qLen); + } + TORCH_CHECK(kvPages.is_cuda() && kvPages.is_contiguous() && kvPages.dim() == 4 + && kvPages.scalar_type() == at::kFloat8_e4m3fn && kvPages.size(1) == numHeadsKv + && kvPages.size(2) == tokensPerBlock && kvPages.size(3) == drk::kHeadDim, + "kv_pages must be the E4M3 [sub-pages, Hkv, tokens_per_block, 128] cache"); + TORCH_CHECK(blockTables.is_cuda() && blockTables.scalar_type() == at::kInt && blockTables.dim() == 3 + && blockTables.size(0) >= requests && blockTables.size(1) == 2 && blockTables.stride(2) == 1, + "block_tables must be the int32 [requests, 2, max blocks] K / V sub-page table"); + TORCH_CHECK(qOut.is_cuda() && qOut.is_contiguous() && qOut.scalar_type() == at::kFloat8_e4m3fn + && qOut.numel() == qkv.size(0) * numHeadsQ * drk::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{qkv.device()}; + drk::Params params{}; + params.qkv = reinterpret_cast<__nv_bfloat16 const*>(qkv.data_ptr()); + params.cosSin = reinterpret_cast(cosSin.data_ptr()); + params.seqLens = seqLens.data_ptr(); + params.specPositionOffsets = specPositionOffsets.has_value() ? specPositionOffsets->data_ptr() : nullptr; + params.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + params.pool = reinterpret_cast<__nv_fp8_e4m3*>(kvPages.data_ptr()); + params.blockTable = blockTables.data_ptr(); + params.blockTableRequestStride = blockTables.stride(0); + params.blockTableKvStride = blockTables.stride(1); + params.numTokens = static_cast(qkv.size(0)); + params.qLen = static_cast(qLen); + params.numHeadsQ = static_cast(numHeadsQ); + params.numHeadsKv = static_cast(numHeadsKv); + params.rotaryDim = static_cast(rotaryDim); + params.tokensPerBlock = static_cast(tokensPerBlock); + params.ropeMode = static_cast(ropeMode); + cudaError_t const err + = drk::launch(params, pdl, maxSharedCarveout, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp2_drafter_rope_fp8_kv: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp2_drafter_rope_fp8_kv(Tensor qkv, Tensor cos_sin, Tensor seq_lens, " + "Tensor? spec_position_offsets, Tensor(a!) kv_pages, Tensor block_tables, Tensor(b!) q_out, int q_len, " + "int num_heads_q, int num_heads_kv, int rotary_dim, int tokens_per_block, int rope_mode, bool pdl, bool " + "max_shared_carveout=True) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp2_drafter_rope_fp8_kv", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2DrafterRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp new file mode 100644 index 000000000000..de42edcc82c0 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxWidth = 32; + +void check2d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 2 && t.size(1) <= kMaxWidth, + "minimax_m3_sm107_tp2_fold_tail: ", name, " must be a 2-D int32 CUDA view at most 32 wide"); +} + +void check1d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp2_fold_tail: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +} // namespace + +//! The speculative step's forward-tail store and next-step staging in one launch (the C++ form of the Triton +//! kernels/m3_fold_tail.py, the same writes). The caller checks the row counts against the buffers (fold_tail). +void minimaxM3Sm107Tp2FoldTail(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + TORCH_CHECK(ctrl.is_cuda() && ctrl.scalar_type() == torch::kInt32 && ctrl.numel() >= 1, + "minimax_m3_sm107_tp2_fold_tail: ctrl must be an int32 CUDA tensor"); + TORCH_CHECK(slots.is_cuda() && slots.scalar_type() == torch::kInt64 && slots.dim() == 1 && slots.stride(0) == 1, + "minimax_m3_sm107_tp2_fold_tail: slots must be a contiguous 1-D int64 CUDA tensor"); + TORCH_CHECK(rows > 0 && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, + "minimax_m3_sm107_tp2_fold_tail: rows must be positive and tokens_per_step in [1, 32]"); + check2d(srcNew, "src_new"); + check2d(dstNew, "dst_new"); + check2d(srcNext, "src_next"); + check2d(dstNext, "dst_next"); + check2d(srcDraft, "src_draft"); + check2d(dstDraft, "dst_draft"); + check1d(srcLens, "src_lens"); + check1d(dstLens, "dst_lens"); + check1d(ids, "ids"); + check1d(draftIds, "draft_ids"); + check1d(posOff, "pos_off"); + check1d(kvOff, "kv_off"); + c10::cuda::CUDAGuard const deviceGuard{ctrl.device()}; + + tensorrt_llm::kernels::minimaxM3Sm107Tp2::FoldTailParams p{}; + p.ctrl = ctrl.data_ptr(); + p.slots = slots.data_ptr(); + p.srcNew = srcNew.data_ptr(); + p.sn0 = srcNew.stride(0); + p.sn1 = srcNew.stride(1); + p.snW = static_cast(srcNew.size(1)); + p.dstNew = dstNew.data_ptr(); + p.dn0 = dstNew.stride(0); + p.dn1 = dstNew.stride(1); + p.dnW = static_cast(dstNew.size(1)); + p.srcNext = srcNext.data_ptr(); + p.sx0 = srcNext.stride(0); + p.sx1 = srcNext.stride(1); + p.sxW = static_cast(srcNext.size(1)); + p.dstNext = dstNext.data_ptr(); + p.dx0 = dstNext.stride(0); + p.dx1 = dstNext.stride(1); + p.dxW = static_cast(dstNext.size(1)); + p.srcDraft = srcDraft.data_ptr(); + p.sd0 = srcDraft.stride(0); + p.sd1 = srcDraft.stride(1); + p.sdW = static_cast(srcDraft.size(1)); + p.dstDraft = dstDraft.data_ptr(); + p.dd0 = dstDraft.stride(0); + p.dd1 = dstDraft.stride(1); + p.ddW = static_cast(dstDraft.size(1)); + p.srcLens = srcLens.data_ptr(); + p.dstLens = dstLens.data_ptr(); + p.ids = ids.data_ptr(); + p.draftIds = draftIds.data_ptr(); + p.posOff = posOff.data_ptr(); + p.kvOff = kvOff.data_ptr(); + p.tokensPerStep = static_cast(tokensPerStep); + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeFoldTail( + p, static_cast(rows), at::cuda::getCurrentCUDAStream(ctrl.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_fold_tail launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_fold_tail(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) dst_new, " + "Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_fold_tail", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2FoldTail); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp new file mode 100644 index 000000000000..8db408008dc2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107Tp2::fullCtxReduce; + +//! Merges the PrimTS split-KV decode scratch (partial_o [B, Hkv, S, R, 128] 16-bit, partial_stats [B, Hkv, S, R] +//! fp32 log2-LSE) of the MiniMax-M3 TP2 full-context decode into out (BF16 [B * SQ, Hkv * hR, 128]). +void minimaxM3Sm107Tp2FullCtxReduce(torch::Tensor const& partialO, torch::Tensor const& partialStats, + torch::Tensor const& seqLensKv, torch::Tensor& out, int64_t seqLenQ, int64_t qTokensPerCta, int64_t tileKv, + int64_t instsKv, bool causal, bool pdl) +{ + TORCH_CHECK(partialO.is_cuda() && partialO.is_contiguous() && partialO.dim() == 5); + TORCH_CHECK(partialO.scalar_type() == at::kHalf || partialO.scalar_type() == at::kBFloat16); + TORCH_CHECK(partialO.size(4) == fcr::kHeadDim && partialO.size(2) <= fcr::kMaxSplits); + TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() + && partialStats.sizes() == partialO.sizes().slice(0, 4)); + TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + int const rows = static_cast(partialO.size(3)); + TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); + int const headsPerKv = rows / static_cast(seqLenQ); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.is_contiguous() + && out.numel() == partialO.size(0) * seqLenQ * partialO.size(1) * headsPerKv * fcr::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{partialO.device()}; + fcr::Params params{partialO.data_ptr(), partialStats.data_ptr(), seqLensKv.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), static_cast(partialO.size(0)), + static_cast(partialO.size(1)), static_cast(partialO.size(2)), rows, static_cast(seqLenQ), + headsPerKv, static_cast(qTokensPerCta), static_cast(tileKv), static_cast(instsKv), causal, + partialO.scalar_type() == at::kBFloat16}; + cudaError_t const err + = fcr::invokeFullCtxReduce(params, pdl, at::cuda::getCurrentCUDAStream(partialO.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_fullctx_reduce: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_fullctx_reduce(Tensor partial_o, Tensor partial_stats, Tensor seq_lens_kv, " + "Tensor(a!) out, int seq_len_q, int q_tokens_per_cta, int tile_kv, int insts_kv, bool causal, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_fullctx_reduce", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2FullCtxReduce); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp new file mode 100644 index 000000000000..a0f27263a152 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp @@ -0,0 +1,181 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h" + +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxTokens = 48; +constexpr int64_t kCluster = 8; +constexpr int64_t kMaxWorld = 8; + +} // namespace + +//! MiniMax-M3 SM107 TP2 multicast all-reduce (stock rounding points: BF16 rank sum, then the residual add) + residual + +//! RMSNorm (+ optional MoE finalize, + optional MXFP8 output). +//! +//! @param x BF16 [M, H] partial, or with `perm` the permuted FC2 output [rows, H]. +//! @param x2 Optional BF16 [M, H] second addend (the shared expert in finalize mode). +//! @param residual BF16 [M, H], 1 <= M <= 48. +//! @param normWeight BF16 [H]. +//! @param wsUc This rank's unicast view of the exchange workspace (int32 words); re-armed in place. +//! @param wsMcPtr Device address of the workspace's multicast mapping. +//! @param flags Int32 per-CTA call-parity counters, at least maxt * 8; updated in place. +//! @param cluster Cluster size; 8 is the only one built. +//! @param perm, scales Finalize mode: int32 [M, topk] permuted-row index and fp32 / bf16 [M, topk] routing weights. +//! @param maxt Token capacity of the workspace bucket (8, 16 or 48). +//! @param quant Also return the MXFP8 copy of the normed rows (E4M3 [M, H], swizzled UE8M0 scales). +//! @param prepushed 0, or 2 / 3 when producers already pushed this rank's rows (see the kernel header). +//! @param opts Epilogue variant; every variant is bit-identical. +//! @param capOut Optional BF16 [>= M, H] view (row stride a multiple of 8, 16-byte aligned) that receives the Eagle3 +//! capture of the updated residual: m3_spec_tail capture_norm with capW [H] and capEps, bit for bit (H 6144; not the +//! drafter boundary, opts bit 16). +//! @return [normed, updated residual] or, with `quant`, [normed, updated residual, E4M3 rows, scale bytes]. +std::vector minimaxM3Sm107Tp2McAllreduceNorm(torch::Tensor const& x, + std::optional const& x2, torch::Tensor const& residual, torch::Tensor const& normWeight, double eps, + torch::Tensor& wsUc, int64_t wsMcPtr, torch::Tensor& flags, int64_t rank, int64_t world, int64_t cluster, bool pdl, + std::optional const& perm, std::optional const& scales, int64_t maxt, bool quant, + int64_t prepushed, int64_t opts, std::optional const& capOut, + std::optional const& capW, double capEps) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 107, "minimax_m3_sm107_tp2_mc_allreduce_norm requires SM107, got SM", sm); + TORCH_CHECK(cluster == kCluster, "minimax_m3_sm107_tp2_mc_allreduce_norm: only cluster 8 is built, got ", cluster); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107Tp2McAllreduceNormSupported( + static_cast(maxt), static_cast(opts)), + "minimax_m3_sm107_tp2_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); + TORCH_CHECK( + world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); + + tensorrt_llm::kernels::MinimaxM3Sm107Tp2McAllreduceNormParams p; + if (x2.has_value()) + { + TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), + "x2 must be a contiguous bf16 tensor shaped like residual"); + p.x2 = reinterpret_cast<__nv_bfloat16 const*>(x2->data_ptr()); + } + if (perm.has_value()) + { + // Finalize mode: x is the permuted FC2 output [rows, H]. + TORCH_CHECK(scales.has_value() && perm->scalar_type() == at::kInt && perm->is_contiguous() && perm->dim() == 2 + && perm->size(0) == residual.size(0), + "perm must be a contiguous int32 [M, topk] tensor, with scales"); + TORCH_CHECK(scales->is_contiguous() && scales->numel() == perm->numel() + && (scales->scalar_type() == at::kFloat || scales->scalar_type() == at::kBFloat16), + "scales must be a contiguous fp32 or bf16 tensor with one weight per perm entry"); + TORCH_CHECK(x.dim() == 2 && x.size(1) == residual.size(1), "fc2 must be [rows, H]"); + p.fc2 = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.perm = perm->data_ptr(); + p.scales = scales->data_ptr(); + p.scaleFp32 = scales->scalar_type() == at::kFloat; + p.topk = static_cast(perm->size(1)); + } + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2, + "x must be a contiguous 2-D bf16 CUDA tensor"); + TORCH_CHECK(residual.dim() == 2 && residual.size(0) >= 1 && residual.size(0) <= std::min(kMaxTokens, maxt), + "residual must hold 1..min(48, bucket) rows"); + TORCH_CHECK(residual.scalar_type() == at::kBFloat16 && residual.is_contiguous() + && (perm.has_value() || residual.sizes() == x.sizes()), + "residual must be a contiguous bf16 tensor shaped like x"); + TORCH_CHECK( + normWeight.scalar_type() == at::kBFloat16 && normWeight.numel() == residual.size(1), "norm_w must be bf16 [H]"); + TORCH_CHECK(wsUc.scalar_type() == at::kInt && wsUc.is_cuda(), "ws_uc must be an int32 CUDA tensor"); + TORCH_CHECK(flags.scalar_type() == at::kInt && flags.is_contiguous() && flags.numel() >= maxt * kCluster, + "flags must be a contiguous int32 tensor of at least bucket * 8 counters"); + + int64_t const m = residual.size(0), h = residual.size(1); + TORCH_CHECK(!quant || h % 128 == 0, "MXFP8 output needs H % 128 == 0"); + c10::cuda::CUDAGuard const deviceGuard{residual.device()}; + auto outNormed = torch::empty_like(residual); + auto outResidual = torch::empty_like(residual); + auto outQ = quant ? torch::empty({m, h}, residual.options().dtype(at::kFloat8_e4m3fn)) : torch::Tensor(); + auto outSf = quant ? torch::empty({(m + 127) / 128 * 128 * (h / 32)}, residual.options().dtype(at::kByte)) + : torch::Tensor(); + + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.numTokens = static_cast(m); + p.hidden = static_cast(h); + p.residual = reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()); + p.normWeight = reinterpret_cast<__nv_bfloat16 const*>(normWeight.data_ptr()); + p.eps = static_cast(eps); + p.wsUc = wsUc.data_ptr(); + p.wsMc = reinterpret_cast(wsMcPtr); + p.flags = flags.data_ptr(); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.outNormed = reinterpret_cast<__nv_bfloat16*>(outNormed.data_ptr()); + p.outResidual = reinterpret_cast<__nv_bfloat16*>(outResidual.data_ptr()); + p.outQuant = quant ? reinterpret_cast(outQ.data_ptr()) : nullptr; + p.outSf = quant ? outSf.data_ptr() : nullptr; + p.prepushed = static_cast(prepushed); + if (capOut.has_value()) + { + TORCH_CHECK((opts & 16) == 0, "the capture epilogue is not built for the drafter boundary (opts bit 16)"); + TORCH_CHECK(capOut->scalar_type() == at::kBFloat16 && capOut->dim() == 2 && capOut->size(0) >= m + && capOut->size(1) == h && capOut->stride(1) == 1 && capOut->stride(0) % 8 == 0 + && reinterpret_cast(capOut->data_ptr()) % 16 == 0 && capOut->device() == residual.device(), + "cap_out must be a bf16 [>= M, H] view with unit column stride, a row stride that is a multiple of 8 and " + "16-byte alignment"); + TORCH_CHECK( + capW.has_value() && capW->scalar_type() == at::kBFloat16 && capW->is_contiguous() && capW->numel() == h, + "cap_w must be a contiguous bf16 [H] tensor"); + p.capOut = reinterpret_cast<__nv_bfloat16*>(capOut->data_ptr()); + p.capStride = capOut->stride(0); + p.capWeight = reinterpret_cast<__nv_bfloat16 const*>(capW->data_ptr()); + p.capEps = static_cast(capEps); + } + tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp2McAllreduceNorm(p, static_cast(maxt), + static_cast(opts), pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); + if (quant) + { + return {outNormed, outResidual, outQ, outSf}; + } + return {outNormed, outResidual}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_mc_allreduce_norm(Tensor x, Tensor? x2, Tensor residual, Tensor norm_w, float eps, " + "Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) flags, int rank, int world, int cluster, bool pdl, " + "Tensor? perm, Tensor? scales, int maxt, bool quant, int prepushed, int opts, Tensor(c!)? cap_out=None, " + "Tensor? cap_w=None, float cap_eps=0.) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_mc_allreduce_norm", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2McAllreduceNorm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp new file mode 100644 index 000000000000..ed50e6d10ba2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp @@ -0,0 +1,185 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp2::producerPw2; + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPageSize = 128; + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "sm107_tp2_producer_pw: ", name, " must be a CUDA tensor"); +} + +} // namespace + +//! The SM107 TP2 NVFP4 QKV + indexer norm/RoPE/KV-insert producer (same arguments and outputs as +//! minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert at 32 query, 2 KV and 2 index heads, head_dim 128, rotary_dim 64): +//! returns (q FP8 [T, 32, 128], index_q FP8 [T, 2, 128]) and writes the NVFP4 K/V of both KV heads and the FP8 index-K +//! of each token's cache slot. `early` triggers the dependents on entry; `prewait_step_inputs` reads the position ids, +//! the slots and the RoPE rows before the grid-dependency wait (none of them may be written by a PDL predecessor). +std::vector minimaxM3Sm107Tp2ProducerPw(torch::Tensor const& packed, torch::Tensor& kvData, + torch::Tensor& kvScale, torch::Tensor& indexK, torch::Tensor const& outCacheLoc, torch::Tensor const& kvQuantScale, + double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds, bool pdl, + bool early, bool prewaitStepInputs) +{ + checkCuda(packed, "packed"); + for (torch::Tensor const* t : std::initializer_list{&kvData, &kvScale, &indexK, &outCacheLoc, + &kvQuantScale, &qWeight, &kWeight, &indexQWeight, &indexKWeight, &rotaryCosSin, &positionIds}) + { + TORCH_CHECK(t->is_cuda() && t->get_device() == packed.get_device(), + "sm107_tp2_producer_pw: every operand must be on packed's CUDA device"); + } + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0 + && reinterpret_cast(kvData.data_ptr()) % 2 == 0 + && reinterpret_cast(indexK.data_ptr()) % 4 == 0, + "sm107_tp2_producer_pw: packed and the caches must start at the vector-load / store alignment"); + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp2_producer_pw requires SM100 or newer"); + int64_t const headsQ = pw::kHeadsQ; + int64_t const headsKV = pw::kHeadsKV; + int64_t const headsIndex = pw::kHeadsIndex; + int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; + TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim + && packed.is_contiguous(), + "sm107_tp2_producer_pw: packed must be contiguous BF16 [T, 39 * 128]"); + // The kernel's layout (the native producer's checkMinimaxM3Nvfp4HndKVPools at two KV heads): HND pools + // [pages, 2 planes, 2 heads, 128 tokens, 64 data / 8 scale bytes], dense token rows and head blocks, planes and + // pages that do not overlap. + TORCH_CHECK(kvData.scalar_type() == at::kByte && kvData.dim() == 5 && kvData.size(0) > 0 && kvData.size(1) == 2 + && kvData.size(2) == headsKV && kvData.size(3) == kPageSize && kvData.size(4) == kHeadDim / 2, + "sm107_tp2_producer_pw: kv_data must be the NVFP4 pool [pages, 2, 2, 128, 64] uint8"); + TORCH_CHECK(kvScale.scalar_type() == at::kByte && kvScale.dim() == 5 && kvScale.size(0) == kvData.size(0) + && kvScale.size(1) == 2 && kvScale.size(2) == headsKV && kvScale.size(3) == kPageSize + && kvScale.size(4) == kHeadDim / 16, + "sm107_tp2_producer_pw: kv_scale must be the E4M3 scale pool [pages, 2, 2, 128, 8] uint8"); + TORCH_CHECK(kvData.stride(4) == 1 && kvData.stride(3) == kvData.size(4) + && kvData.stride(2) == kvData.size(3) * kvData.stride(3), + "sm107_tp2_producer_pw: kv_data must have dense token rows and head blocks"); + TORCH_CHECK(kvScale.stride(4) == 1 && kvScale.stride(3) == kvScale.size(4) + && kvScale.stride(2) == kvScale.size(3) * kvScale.stride(3), + "sm107_tp2_producer_pw: kv_scale must have dense token rows and head blocks"); + TORCH_CHECK(kvData.stride(1) >= kvData.size(2) * kvData.stride(2) && kvData.stride(0) >= 2 * kvData.stride(1) + && kvScale.stride(1) >= kvScale.size(2) * kvScale.stride(2) && kvScale.stride(0) >= 2 * kvScale.stride(1), + "sm107_tp2_producer_pw: pool planes or pages overlap"); + TORCH_CHECK(kvData.stride(0) % 2 == 0 && kvData.stride(1) % 2 == 0, + "sm107_tp2_producer_pw: kv_data page and plane strides must keep the uint16 stores aligned"); + TORCH_CHECK(indexK.scalar_type() == at::kFloat8_e4m3fn && indexK.dim() == 4 && indexK.size(0) == kvData.size(0) + && indexK.size(1) == 1 && indexK.size(2) == kPageSize && indexK.size(3) == kHeadDim && indexK.stride(3) == 1 + && indexK.stride(2) == kHeadDim && indexK.stride(0) >= kPageSize * kHeadDim && indexK.stride(0) % 4 == 0, + "sm107_tp2_producer_pw: index_k must be FP8 [pages, 1, 128, 128] with a page stride keeping the 32-bit stores " + "aligned"); + int64_t const numTokens = packed.size(0); + TORCH_CHECK( + outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous() && outCacheLoc.numel() >= numTokens, + "sm107_tp2_producer_pw: out_cache_loc must be int32 [>= T]"); + TORCH_CHECK( + positionIds.scalar_type() == at::kInt && positionIds.is_contiguous() && positionIds.numel() >= numTokens, + "sm107_tp2_producer_pw: position_ids must be int32 [>= T]"); + TORCH_CHECK(kvQuantScale.scalar_type() == at::kFloat && kvQuantScale.is_contiguous() && kvQuantScale.numel() >= 3, + "sm107_tp2_producer_pw: kv_quant_scale must be fp32 [>= 3]"); + TORCH_CHECK(rotaryCosSin.scalar_type() == at::kFloat && rotaryCosSin.is_contiguous() && rotaryCosSin.dim() == 3 + && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == 32 + && reinterpret_cast(rotaryCosSin.data_ptr()) % 16 == 0, + "sm107_tp2_producer_pw: rotary_cos_sin must be fp32 [positions, 2, 32], 16-byte aligned (read as float4)"); + for (auto const* w : {&qWeight, &kWeight, &indexQWeight, &indexKWeight}) + { + TORCH_CHECK(w->scalar_type() == at::kBFloat16 && w->is_contiguous() && w->numel() == kHeadDim + && reinterpret_cast(w->data_ptr()) % 8 == 0, + "sm107_tp2_producer_pw: norm weights must be BF16 [128], 8-byte aligned (read as 4-value words)"); + } + auto const fp8 = packed.options().dtype(at::kFloat8_e4m3fn); + auto qOut = torch::empty({numTokens, headsQ, kHeadDim}, fp8); + auto indexQOut = torch::empty({numTokens, headsIndex, kHeadDim}, fp8); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + pw::Params p{}; + p.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + p.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + p.indexQOut = reinterpret_cast<__nv_fp8_e4m3*>(indexQOut.data_ptr()); + p.kvData = kvData.data_ptr(); + p.kvScale = kvScale.data_ptr(); + p.indexK = reinterpret_cast<__nv_fp8_e4m3*>(indexK.data_ptr()); + p.outCacheLoc = outCacheLoc.data_ptr(); + p.kvQuantScale = kvQuantScale.data_ptr(); + p.dataPageStride = kvData.stride(0); + p.dataPlaneStride = kvData.stride(1); + p.dataHeadStride = kvData.stride(2); + p.dataTokenStride = kvData.stride(3); + p.scalePageStride = kvScale.stride(0); + p.scalePlaneStride = kvScale.stride(1); + p.scaleHeadStride = kvScale.stride(2); + p.indexPageStride = indexK.stride(0); + p.indexTokenStride = indexK.stride(2); + p.numPages = kvData.size(0); + p.numTokens = static_cast(numTokens); + p.eps = static_cast(eps); + p.qWeight = reinterpret_cast<__nv_bfloat16 const*>(qWeight.data_ptr()); + p.kWeight = reinterpret_cast<__nv_bfloat16 const*>(kWeight.data_ptr()); + p.indexQWeight = reinterpret_cast<__nv_bfloat16 const*>(indexQWeight.data_ptr()); + p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); + p.rotaryCosSin = rotaryCosSin.data_ptr(); + p.positionIds = positionIds.data_ptr(); + cudaError_t const err = pw::invokeProducerPw2(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "sm107_tp2_producer_pw launch: ", cudaGetErrorString(err)); + } + return {qOut, indexQOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_producer_pw(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) index_k, " + "Tensor out_cache_loc, Tensor kv_quant_scale, float eps, Tensor q_weight, Tensor k_weight, " + "Tensor index_q_weight, Tensor index_k_weight, Tensor rotary_cos_sin, Tensor position_ids, bool pdl, " + "bool early, bool prewait_step_inputs=False) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_producer_pw", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2ProducerPw); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp new file mode 100644 index 000000000000..415042d5c082 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp @@ -0,0 +1,110 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3Sm107Tp2::sampledAccept; + +//! uint64 words of the sampled-accept exchange workspace for `world` ranks. +int64_t minimaxM3Sm107Tp2SampledAcceptWsWords(int64_t world) +{ + return sa::sampledAcceptWsWords(static_cast(world)); +} + +//! Sampled (temperature + top-p) verification from this rank's BF16 logits shard [rows, Vs], rows = nreq * +//! (draft_len + 1) <= 48: the sampled token of every row (tokens), then the strict acceptance against draft (accTokens, +//! numAcc, gatherIds, optional nextNew column 0). Per-row Philox (seeds, offsets); the exchange goes through the +//! multicast workspace wsMcPtr / wsUc with the (finished rows, call) counters. +void minimaxM3Sm107Tp2SampledAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor const& temperatures, + torch::Tensor const& topPs, torch::Tensor const& seeds, torch::Tensor const& offsets, torch::Tensor& wsUc, + int64_t wsMcPtr, torch::Tensor& counters, torch::Tensor& tokens, int64_t rank, int64_t world, + torch::Tensor const& draft, torch::Tensor const& batchIdx, torch::Tensor& accTokens, torch::Tensor& numAcc, + torch::Tensor& gatherIds, int64_t forceTotal, std::optional const& nextNew, + std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, + bool prewaitParams, bool earlyTrigger, bool ilpWalk) +{ + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + sa::Args a; + a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + a.ld = logits.stride(0); + a.rows = static_cast(logits.size(0)); + a.Vs = static_cast(logits.size(1)); + TORCH_CHECK(a.Vs <= sa::kMaxVs && a.rows <= sa::kMaxRows && world <= sa::kMaxWorld + && static_cast(vocabOffset) + a.Vs <= (1LL << sa::kIdBits)); + a.vocab_offset = static_cast(vocabOffset); + a.temperatures = temperatures.data_ptr(); + a.top_ps = topPs.data_ptr(); + a.seeds = reinterpret_cast(seeds.data_ptr()); + a.offsets = reinterpret_cast(offsets.data_ptr()); + a.ws_uc = reinterpret_cast(wsUc.data_ptr()); + a.ws_mc = reinterpret_cast(wsMcPtr); + a.counters = reinterpret_cast(counters.data_ptr()); + a.tokens = tokens.data_ptr(); + a.extra_attempts = extraAttempts.has_value() ? extraAttempts->data_ptr() : nullptr; + a.rank = static_cast(rank); + a.world = static_cast(world); + a.draft = draft.data_ptr(); + a.batch_idx = batchIdx.data_ptr(); + a.acc_tokens = accTokens.data_ptr(); + a.num_acc = numAcc.data_ptr(); + a.gather_ids = reinterpret_cast(gatherIds.data_ptr()); + a.draft_len = static_cast(draft.size(1)); + a.nreq = static_cast(draft.size(0)); + a.force_total = static_cast(forceTotal); + a.next_new = nextNew.has_value() ? nextNew->data_ptr() : nullptr; + a.next_stride = nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0; + a.force_vec = forceVec.has_value() ? forceVec->data_ptr() : nullptr; + a.prewait_params = prewaitParams ? 1 : 0; + a.early_trigger = earlyTrigger ? 1 : 0; + a.ilp_walk = ilpWalk ? 1 : 0; + cudaError_t const err = sa::invokeSampledAccept(a, usePdl, at::cuda::getCurrentCUDAStream(logits.get_device())); + TORCH_CHECK(err == cudaSuccess, "m3_sampled_accept launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// ws_mc_ptr is the multicast address of ws_uc's mapping: the op writes through it, which the schema cannot show. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + // No tensor arguments, so no dispatch key: registered with the schema as a catch-all kernel. + m.def("minimax_m3_sm107tp2_sampled_accept_ws_words(int world) -> int", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SampledAcceptWsWords); + m.def( + "minimax_m3_sm107tp2_sampled_accept(Tensor logits, int vocab_offset, Tensor temperatures, Tensor top_ps, " + "Tensor seeds, Tensor offsets, Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) counters, Tensor(c!) tokens, " + "int rank, int world, Tensor draft, Tensor batch_idx, Tensor(d!) acc_tokens, Tensor(e!) num_acc, " + "Tensor(f!) gather_ids, int force_total, Tensor(g!)? next_new, Tensor? force_vec, " + "Tensor(h!)? extra_attempts, bool use_pdl, bool prewait_params=True, bool early_trigger=False, " + "bool ilp_walk=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp2_sampled_accept", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SampledAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SelectLongOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SelectLongOp.cpp new file mode 100644 index 000000000000..b54249840c93 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SelectLongOp.cpp @@ -0,0 +1,85 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, into a caller-owned table. +//! +//! @param scores Float32 CUDA tensor [num_kv_heads, num_blocks, total_q], any strides. +//! @param nValidBlocks Contiguous int32 CUDA tensor with at least total_q elements. +//! @param out Int32 CUDA tensor of logical shape [total_q, num_kv_heads, 16] with unit stride in the last dim (the +//! head-major backing the target hands over is a permuted view). Every entry is written. +//! @param triggerPhase Where the dependent grid is released without earlyTrigger: 0 at this grid's exit, 1 right +//! after every score of the row is loaded (both after this grid's dependency wait). +void minimaxM3Sm107Tp2SelectLong(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, torch::Tensor& out, + int64_t initBlocks, int64_t localBlocks, bool usePdl, bool earlyTrigger, int64_t triggerPhase) +{ + TORCH_CHECK(scores.is_cuda() && scores.scalar_type() == torch::kFloat32 && scores.dim() == 3, + "minimax_m3_sm107_tp2_select_long: scores must be a 3-D float32 CUDA tensor [num_kv_heads, num_blocks, " + "total_q]"); + TORCH_CHECK(nValidBlocks.is_cuda() && nValidBlocks.scalar_type() == torch::kInt32 && nValidBlocks.is_contiguous() + && nValidBlocks.numel() >= scores.size(2), + "minimax_m3_sm107_tp2_select_long: n_valid_blocks must be contiguous int32 CUDA with at least total_q " + "elements"); + int64_t const numKvHeads = scores.size(0); + int64_t const numBlocks = scores.size(1); + int64_t const totalQ = scores.size(2); + TORCH_CHECK(numBlocks <= tensorrt_llm::kernels::minimaxM3Sm107Tp2::kSelectLong2MaxBlocks, + "minimax_m3_sm107_tp2_select_long supports at most ", + tensorrt_llm::kernels::minimaxM3Sm107Tp2::kSelectLong2MaxBlocks, " blocks, got ", numBlocks); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kInt32 && out.dim() == 3 && out.size(0) == totalQ + && out.size(1) == numKvHeads && out.size(2) == 16 && out.stride(2) == 1, + "minimax_m3_sm107_tp2_select_long: out must be int32 [total_q, num_kv_heads, 16] with unit last-dim stride"); + + c10::cuda::CUDAGuard const deviceGuard{scores.device()}; + auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device()); + cudaError_t const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSelectLong2(scores.data_ptr(), + scores.stride(0), scores.stride(1), scores.stride(2), nValidBlocks.data_ptr(), out.data_ptr(), + out.stride(0), out.stride(1), static_cast(numKvHeads), static_cast(numBlocks), + static_cast(totalQ), static_cast(initBlocks), static_cast(localBlocks), usePdl, + earlyTrigger, static_cast(triggerPhase), stream); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_sm107_tp2_select_long launch failed: ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_select_long(Tensor scores, Tensor n_valid_blocks, Tensor(a!) out, int init_blocks, " + "int local_blocks, bool use_pdl=False, bool early_trigger=False, int trigger_phase=0) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_select_long", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SelectLong); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp new file mode 100644 index 000000000000..52f60de1b7ca --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp2/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, +//! d]. +void minimaxM3Sm107Tp2SiluAndMul(torch::Tensor const& x, torch::Tensor& out) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp2_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp2_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp2_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); + TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp2_silu_and_mul: at most 65535 rows"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err + = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), + reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_silu_and_mul launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SiluAndMul); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp new file mode 100644 index 000000000000..fb8fc9eaa8ee --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeadsPerKv = 16; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kTopK = 16; + +void checkBytes(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, "minimax_m3_sm107_tp2_sparse_decode_lc: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3 long-context sparse decode over the selected NVFP4 blocks, SM107 TP2 geometry (16 query +//! heads per KV head, 2 KV heads per rank, head_dim 128, page 128, top-k 16). +//! +//! @param q E4M3 bytes [T, 16 * H, 128], unit stride in the last dim. +//! @param k, v Packed E2M1 bytes [pages, H, 128, 64] (the per-layer HND views). +//! @param kBlockScale, vBlockScale E4M3 bytes [pages, H, 128, 8]: K token-major, V in 4x4 token-quad order. +//! @param kGlobalScale, vGlobalScale One-element float32 per-tensor dequant scales. +//! @param topk Int32 [H, T, 16] ascending selected block ids with -1 tail (the head-major selector table). +//! @param out BF16 [T, H * 2048], written. q8 / sf, when given, receive the MXFP8 copy and its swizzled UE8M0 scales. +//! @param n1Ids Optional int32 [H, rows, 16] (unit last-dim stride): this layer's selected block ids at its previous +//! decode step, read to prefetch those blocks into L2 before the grid dependency and overwritten with this step's; +//! query rows at or past `rows` skip it. Initialize to -1. A hint: `out` does not depend on it. +void minimaxM3Sm107Tp2SparseDecodeLc(torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& kBlockScale, torch::Tensor const& vBlockScale, torch::Tensor const& kGlobalScale, + torch::Tensor const& vGlobalScale, torch::Tensor const& topk, torch::Tensor const& blockTable, + torch::Tensor const& seqLens, torch::Tensor& out, std::optional q8, std::optional sf, + int64_t numSfKTiles, int64_t decodeQueryLen, double smScale, bool linearSoftmax, bool usePdl, bool prewait, + std::optional n1Ids) +{ + TORCH_CHECK(q.is_cuda(), "minimax_m3_sm107_tp2_sparse_decode_lc: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m3_sm107_tp2_sparse_decode_lc requires an SM100-family GPU"); + checkBytes(q, "q"); + checkBytes(k, "k"); + checkBytes(v, "v"); + checkBytes(kBlockScale, "k_block_scale"); + checkBytes(vBlockScale, "v_block_scale"); + TORCH_CHECK(k.dim() == 4 && k.size(2) == 128 && k.size(3) == kHeadDim / 2 && k.stride(3) == 1 + && v.sizes() == k.sizes() && v.strides() == k.strides(), + "minimax_m3_sm107_tp2_sparse_decode_lc: k and v must be [pages, H, 128, 64] with matching strides"); + int64_t const kvHeads = k.size(1); + TORCH_CHECK(kBlockScale.dim() == 4 && kBlockScale.size(1) == kvHeads && kBlockScale.size(3) == 8 + && kBlockScale.stride(3) == 1 && kBlockScale.stride(2) == 8 && vBlockScale.sizes() == kBlockScale.sizes() + && vBlockScale.strides() == kBlockScale.strides(), + "minimax_m3_sm107_tp2_sparse_decode_lc: block scales must be contiguous [pages, H, 128, 8] per page"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeadsPerKv * kvHeads && q.size(2) == kHeadDim && q.stride(2) == 1, + "minimax_m3_sm107_tp2_sparse_decode_lc: q must be [T, 16 * H, 128] with unit last-dim stride"); + TORCH_CHECK(kGlobalScale.scalar_type() == torch::kFloat32 && vGlobalScale.scalar_type() == torch::kFloat32, + "minimax_m3_sm107_tp2_sparse_decode_lc: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 3 && topk.size(0) == kvHeads + && topk.size(1) == q.size(0) && topk.size(2) == kTopK, + "minimax_m3_sm107_tp2_sparse_decode_lc: topk must be int32 [H, T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_sm107_tp2_sparse_decode_lc: block_table and seq_lens must be int32"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.size(0) == q.size(0) + && out.size(1) == kvHeads * kHeadsPerKv * kHeadDim && out.stride(1) == 1, + "minimax_m3_sm107_tp2_sparse_decode_lc: out must be bfloat16 [T, H * 2048] with unit last-dim stride"); + TORCH_CHECK( + q8.has_value() == sf.has_value(), "minimax_m3_sm107_tp2_sparse_decode_lc: pass both q8 and sf, or neither"); + bool const mxfp8 = q8.has_value(); + if (mxfp8) + { + checkBytes(*q8, "q8"); + checkBytes(*sf, "sf"); + } + + tensorrt_llm::kernels::MinimaxM3Sm107Tp2SparseDecodeLcParams prm{}; + prm.q = reinterpret_cast(q.data_ptr()); + prm.q_tok = q.stride(0); + prm.q_head = q.stride(1); + prm.k = reinterpret_cast(k.data_ptr()); + prm.v = reinterpret_cast(v.data_ptr()); + prm.k_blk = k.stride(0); + prm.v_blk = v.stride(0); + prm.k_h = k.stride(1); + prm.v_h = v.stride(1); + prm.ksc = reinterpret_cast(kBlockScale.data_ptr()); + prm.vsc = reinterpret_cast(vBlockScale.data_ptr()); + prm.ksc_blk = kBlockScale.stride(0); + prm.vsc_blk = vBlockScale.stride(0); + prm.ksc_h = kBlockScale.stride(1); + prm.vsc_h = vBlockScale.stride(1); + prm.kgs = kGlobalScale.data_ptr(); + prm.vgs = vGlobalScale.data_ptr(); + prm.topk = topk.data_ptr(); + prm.t_h = topk.stride(0); + prm.t_tok = topk.stride(1); + prm.t_k = topk.stride(2); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.seq_lens = seqLens.data_ptr(); + prm.dql = static_cast(decodeQueryLen); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + prm.out_stride = out.stride(0); + prm.q8 = mxfp8 ? reinterpret_cast(q8->data_ptr()) : nullptr; + prm.sf = mxfp8 ? reinterpret_cast(sf->data_ptr()) : nullptr; + prm.num_sf_k_tiles = static_cast(numSfKTiles); + prm.num_kv_heads = static_cast(kvHeads); + if (n1Ids.has_value()) + { + TORCH_CHECK(n1Ids->is_cuda() && n1Ids->scalar_type() == torch::kInt32 && n1Ids->dim() == 3 + && n1Ids->size(0) == kvHeads && n1Ids->size(2) == kTopK && n1Ids->stride(2) == 1 + && n1Ids->stride(1) == kTopK, + "minimax_m3_sm107_tp2_sparse_decode_lc: n1_ids must be int32 [H, rows, 16] with dense rows"); + prm.n1_ids = n1Ids->data_ptr(); + prm.n1_h = n1Ids->stride(0); + prm.n1_rows = static_cast(n1Ids->size(1)); + prm.n1_pages = k.size(0); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp2SparseDecodeLc( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, prewait, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_sparse_decode_lc launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_sparse_decode_lc(Tensor q, Tensor k, Tensor v, Tensor k_block_scale, " + "Tensor v_block_scale, Tensor k_global_scale, Tensor v_global_scale, Tensor topk, Tensor block_table, " + "Tensor seq_lens, Tensor(a!) out, Tensor(b!)? q8, Tensor(c!)? sf, int num_sf_k_tiles, int decode_query_len, " + "float sm_scale, bool linear_softmax, bool use_pdl, bool prewait, Tensor(d!)? n1_ids=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_sparse_decode_lc", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SparseDecodeLc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp new file mode 100644 index 000000000000..87b6ae599742 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void check(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp2_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +} // namespace + +//! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch +//! (kernels/minimaxM3Sm107Tp2/specOffsets.h). +void minimaxM3Sm107Tp2SpecOffsets( + torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) +{ + check(pos, "pos"); + check(posOff, "pos_off"); + check(kv, "kv"); + check(kvOff, "kv_off"); + TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), + "minimax_m3_sm107_tp2_spec_offsets: each tensor and its offsets must have the same length"); + if (pos.numel() == 0 && kv.numel() == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{pos.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSpecOffsets(pos.data_ptr(), + posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), + at::cuda::getCurrentCUDAStream(pos.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_spec_offsets launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SpecOffsets); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp new file mode 100644 index 000000000000..9e8229ba9e16 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp2/stepPrologue.h). +//! +//! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the +//! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses +//! the caller guarantees are live. +//! @param epoch Optional one-element int64 step counter, incremented once. +void minimaxM3Sm107Tp2StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) +{ + TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), + "minimax_m3_sm107_tp2_step_prologue: window must be a contiguous uint8 CUDA tensor"); + TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, + "minimax_m3_sm107_tp2_step_prologue: window is shorter than its header and max_desc descriptors"); + c10::cuda::CUDAGuard const deviceGuard{window.device()}; + int64_t* epochPtr = nullptr; + if (epoch.has_value()) + { + TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 + && epoch->device() == window.device(), + "minimax_m3_sm107_tp2_step_prologue: epoch must be a one-element int64 tensor on the window's device"); + epochPtr = epoch->data_ptr(); + } + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeStepPrologue( + window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_step_prologue launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2StepPrologue); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp new file mode 100644 index 000000000000..2839f6036ffc --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp2/swigluOai.h): x bf16 [rows, 2d] -> out bf16 +//! [rows, d], PDL-launched with an early release when `pdl`. +void minimaxM3Sm107Tp2SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp2_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp2_swiglu_oai: x rows must be 4-byte aligned with an even half width"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp2_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSwigluOai( + reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, + at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_swiglu_oai launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SwigluOai); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp new file mode 100644 index 000000000000..d92f2c1cd29b --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp @@ -0,0 +1,128 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace dqk = tensorrt_llm::kernels::minimaxM3Sm107Tp4::denseQkNormRopeFp8; + +//! QK RMSNorm + partial RoPE of a packed BF16 [tokens, (q + k + v) * 128] row of the MiniMax-M3 TP4 dense layers, +//! returned as a new FP8 E4M3 tensor of the same shape (V copy-cast; Q/K rounded to BF16 before E4M3, as stock). +torch::Tensor minimaxM3Sm107Tp4DenseQkNormRopeFp8(torch::Tensor const& qkv, int64_t numHeadsQ, int64_t numHeadsK, + int64_t numHeadsV, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, + double base, bool isNeox, torch::Tensor const& positionIds, bool useGemma) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); + TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); + TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); + TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + at::cuda::CUDAGuard guard(qkv.device()); + auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); + cudaError_t const err + = dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(rotaryDim), + static_cast(eps), qWeight.data_ptr(), kWeight.data_ptr(), static_cast(base), !isNeox, + positionIds.data_ptr(), useGemma, nullptr, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "dense_qk_norm_rope_fp8 launch: ", cudaGetErrorString(err)); + return out; +} + +//! As minimaxM3Sm107Tp4DenseQkNormRopeFp8, and the first numLiveTokens rows' K and V are also stored into the layer's +//! FP8 pool kvBuffers (the HND view [pages, 2, kv heads, tokens per block, 128]) at their int32 cache slots; with +//! qCompact, every row's Q is also stored there compact ([tokens, q heads, 128]). +torch::Tensor minimaxM3Sm107Tp4DenseQkNormRopeFp8Kv(torch::Tensor const& qkv, int64_t numHeadsQ, int64_t numHeadsK, + int64_t numHeadsV, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, + double base, bool isNeox, torch::Tensor const& positionIds, bool useGemma, torch::Tensor const& kvBuffers, + torch::Tensor const& outCacheLoc, int64_t numLiveTokens, std::optional const& qCompact) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); + TORCH_CHECK(numHeadsK == numHeadsV, "the pool keeps one head count for K and V"); + TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); + TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); + TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + TORCH_CHECK(kvBuffers.is_cuda() && kvBuffers.scalar_type() == at::kFloat8_e4m3fn && kvBuffers.dim() == 5); + TORCH_CHECK(kvBuffers.size(1) == 2 && kvBuffers.size(2) == numHeadsK && kvBuffers.size(4) == dqk::kHeadDim + && kvBuffers.stride(4) == 1, + "kv_buffers must be the HND pool view [pages, 2, kv heads, tokens per block, 128]"); + TORCH_CHECK(outCacheLoc.is_cuda() && outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous()); + TORCH_CHECK(numLiveTokens >= 0 && numLiveTokens <= qkv.size(0) && numLiveTokens <= outCacheLoc.numel()); + at::cuda::CUDAGuard guard(qkv.device()); + auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); + dqk::KvCacheView view; + view.base = static_cast<__nv_fp8_e4m3*>(kvBuffers.data_ptr()); + view.outCacheLoc = outCacheLoc.data_ptr(); + view.pageStride = kvBuffers.stride(0); + view.kvStride = kvBuffers.stride(1); + view.headStride = kvBuffers.stride(2); + view.tokenStride = kvBuffers.stride(3); + view.tokensPerBlock = static_cast(kvBuffers.size(3)); + view.numLiveTokens = static_cast(numLiveTokens); + if (qCompact.has_value()) + { + TORCH_CHECK(qCompact->is_cuda() && qCompact->is_contiguous() && qCompact->scalar_type() == at::kFloat8_e4m3fn + && qCompact->numel() == qkv.size(0) * numHeadsQ * dqk::kHeadDim, + "q_compact must be a contiguous FP8 [tokens, q heads, 128] tensor"); + view.qCompact = static_cast<__nv_fp8_e4m3*>(qCompact->data_ptr()); + } + cudaError_t const err + = dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(rotaryDim), + static_cast(eps), qWeight.data_ptr(), kWeight.data_ptr(), static_cast(base), !isNeox, + positionIds.data_ptr(), useGemma, &view, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "dense_qk_norm_rope_fp8_kv launch: ", cudaGetErrorString(err)); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, " + "bool use_gemma) -> Tensor"); + m.def( + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, " + "bool use_gemma, Tensor(a!) kv_buffers, Tensor out_cache_loc, int num_live_tokens, Tensor(b!)? q_compact=None) " + "-> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_dense_qk_norm_rope_fp8", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4DenseQkNormRopeFp8); + m.impl("minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4DenseQkNormRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DraftStepOp.cpp new file mode 100644 index 000000000000..b0a49e8a29fb --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DraftStepOp.cpp @@ -0,0 +1,215 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h" +#include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h" + +#include +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace ds = tensorrt_llm::kernels::minimaxM3Sm107Tp4::draftStep; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_sm107tp4_draft_step expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107tp4_draft_step requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +void checkInt(torch::Tensor const& t, int64_t n, char const* name) +{ + TORCH_CHECK(t.scalar_type() == at::kInt && t.is_contiguous() && t.numel() >= n, "draft_step: ", name, + " must be contiguous int32 with at least ", n, " elements"); +} + +} // namespace + +//! One fused greedy draft pass over `rows` <= 8 requests: the BF16 head (this rank's fragment-ordered shard wf +//! [N / 16, K / 16, 256], or with excOff / excPos / excW / expBase its 12-bit coding, wf uint8 [N * K * 3 / 2]; +//! global ids vocabOffset..) over x[xRows[b]] (or x[b]), the global argmax across the TP ranks through the multicast +//! slots wsMcPtr / wsUc, then draftOut[b, step] = idsNext[b] = token, posNext[b] = posSrc[xRows[b] or b] + 1 and the +//! KV-length bookkeeping (see draftStep.h). +void minimaxM3Sm107Tp4DraftStep(torch::Tensor const& x, OptTensor const& xRows, int64_t rows, torch::Tensor const& wf, + int64_t vocabOffset, torch::Tensor& keys, torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, + int64_t rank, int64_t world, torch::Tensor& draftOut, int64_t step, torch::Tensor& idsNext, + torch::Tensor const& posSrc, torch::Tensor& posNext, torch::Tensor& kvLens, OptTensor const& numAcc, + OptTensor const& seqLens, int64_t draftLen, int64_t numCtas, bool pdl, OptTensor const& outKeys, + OptTensor const& excOff, OptTensor const& excPos, OptTensor const& excW, int64_t expBase, + std::optional> fold, int64_t foldRows, int64_t foldTokensPerStep, bool fixedOrder) +{ + checkArch(x); + for (torch::Tensor const* t : std::initializer_list{ + &wf, &keys, &counters, &wsUc, &draftOut, &idsNext, &posSrc, &posNext, &kvLens}) + { + TORCH_CHECK(t->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + for (OptTensor const* t : + std::initializer_list{&xRows, &numAcc, &seqLens, &outKeys, &excOff, &excPos, &excW}) + { + TORCH_CHECK(!t->has_value() || (*t)->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "draft_step: x must be contiguous BF16 [rows, K]"); + int const K = static_cast(x.size(1)); + TORCH_CHECK(K % 16 == 0, "draft_step: K must be a multiple of 16"); + TORCH_CHECK(rows >= 1 && rows <= ds::kMaxTokens, "draft_step: 1..", ds::kMaxTokens, " rows"); + bool const coded = excOff.has_value(); + TORCH_CHECK(coded == excPos.has_value() && coded == excW.has_value(), + "draft_step: exc_off, exc_pos and exc_w come together"); + int nTiles = 0; + if (coded) + { + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "draft_step: a coded head needs K a multiple of 64"); + TORCH_CHECK(excOff->scalar_type() == at::kInt && excOff->is_contiguous() && excOff->numel() >= 2, + "draft_step: exc_off must be int32 [tiles + 1]"); + nTiles = static_cast(excOff->numel() - 1); + TORCH_CHECK(wf.scalar_type() == at::kByte && wf.is_contiguous() + && wf.numel() == static_cast(nTiles) * 16 * K * 3 / 2, + "draft_step: a coded wf must be the contiguous uint8 stream of [tiles * 16, K]"); + TORCH_CHECK(excPos->scalar_type() == at::kInt && excPos->is_contiguous(), "draft_step: exc_pos must be int32"); + TORCH_CHECK(excW->scalar_type() == at::kFloat && excW->is_contiguous() && excW->numel() == 2 * excPos->numel(), + "draft_step: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "draft_step: exponent base out of range"); + } + else + { + TORCH_CHECK(wf.scalar_type() == at::kBFloat16 && wf.dim() == 3 && wf.is_contiguous() && wf.size(1) == K / 16 + && wf.size(2) == 256, + "draft_step: wf must be the fragment-ordered BF16 weight [N / 16, K / 16, 256]"); + nTiles = static_cast(wf.size(0)); + } + TORCH_CHECK(nTiles >= 1 && numCtas >= 1 && (nTiles + numCtas - 1) / numCtas <= ds::kMaxTilesPerCta, + "draft_step: too many weight tiles per CTA"); + TORCH_CHECK(vocabOffset >= 0 && vocabOffset + 16LL * nTiles <= 0xFFFFFFFFLL, "draft_step: vocab offset"); + if (xRows.has_value()) + { + TORCH_CHECK(xRows->scalar_type() == at::kLong && xRows->is_contiguous() && xRows->numel() >= rows, + "draft_step: x_rows must be contiguous int64 [rows]"); + } + else + { + TORCH_CHECK(x.size(0) >= rows, "draft_step: x must hold a row per request"); + } + TORCH_CHECK(world >= 1 && rank >= 0 && rank < world, "draft_step: rank / world"); + TORCH_CHECK(keys.scalar_type() == at::kLong && keys.numel() >= rows, "draft_step: keys must be int64 [>= rows]"); + TORCH_CHECK(counters.scalar_type() == at::kInt && counters.numel() >= 2, "draft_step: counters must be int32 [2]"); + TORCH_CHECK(wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, + "draft_step: ws_uc must hold [2][8][world] slots"); + TORCH_CHECK(draftOut.scalar_type() == at::kInt && draftOut.dim() == 2 && draftOut.stride(1) == 1 + && draftOut.size(0) >= rows && step >= 0 && step < draftOut.size(1), + "draft_step: draft_out must be int32 [rows, > step] with unit column stride"); + checkInt(idsNext, rows, "ids_next"); + checkInt(posNext, rows, "pos_next"); + checkInt(kvLens, rows, "kv_lens"); + TORCH_CHECK(posSrc.scalar_type() == at::kInt && posSrc.is_contiguous(), "draft_step: pos_src must be int32"); + TORCH_CHECK(numAcc.has_value() == seqLens.has_value(), "draft_step: num_acc and seq_lens come together"); + TORCH_CHECK(!numAcc.has_value() || step == 0, "draft_step: the KV rollback (num_acc) is pass 0's"); + if (numAcc.has_value()) + { + checkInt(*numAcc, rows, "num_acc"); + checkInt(*seqLens, rows, "seq_lens"); + } + if (outKeys.has_value()) + { + TORCH_CHECK(outKeys->scalar_type() == at::kLong && outKeys->is_contiguous() && outKeys->numel() >= rows, + "draft_step: out_keys must be contiguous int64 [rows]"); + } + ds::Params p{}; + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_rows = xRows.has_value() ? reinterpret_cast(xRows->data_ptr()) : nullptr; + p.M = static_cast(rows); + p.wf = reinterpret_cast(wf.data_ptr()); + p.n_tiles = nTiles; + p.vocab_offset = static_cast(vocabOffset); + p.keys = reinterpret_cast(keys.data_ptr()); + p.counters = reinterpret_cast(counters.data_ptr()); + p.ws_uc = reinterpret_cast(wsUc.data_ptr()); + p.ws_mc = reinterpret_cast(wsMcPtr); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.out_keys = outKeys.has_value() ? reinterpret_cast(outKeys->data_ptr()) : nullptr; + p.st = ds::StepArgs{draftOut.data_ptr(), static_cast(draftOut.stride(0)), static_cast(step), + idsNext.data_ptr(), posSrc.data_ptr(), posNext.data_ptr(), kvLens.data_ptr(), + numAcc.has_value() ? numAcc->data_ptr() : nullptr, + seqLens.has_value() ? seqLens->data_ptr() : nullptr, static_cast(draftLen)}; + p.cd = ds::CodedArgs{coded ? excOff->data_ptr() : nullptr, + coded ? reinterpret_cast(excPos->data_ptr()) : nullptr, + coded ? reinterpret_cast(excW->data_ptr()) : nullptr, + static_cast((expBase << 7) | (expBase << 23))}; + if (fold.has_value()) + { + // The 14 fold-tail tensors in buildParams' order (ctrl, slots, src_new, dst_new, src_next, dst_next, + // src_draft, dst_draft, src_lens, dst_lens, ids, draft_ids, pos_off, kv_off): the in-kernel tail's operands, + // packed exactly as the standalone fold tail packs them. + std::vector f = *fold; + TORCH_CHECK(f.size() == 14, "minimax_m3_sm107tp4_draft_step: fold takes 14 tensors, got ", f.size()); + p.fold = minimaxM3Sm107Tp4FoldTailParams::buildParams(f[0], f[1], foldRows, f[2], f[3], f[4], f[5], f[6], f[7], + f[8], f[9], f[10], f[11], f[12], f[13], foldTokensPerStep); + p.hasFold = true; + } + p.fixedOrder = fixedOrder; + check(ds::invokeDraftStep(p, K, static_cast(numCtas), pdl, at::cuda::getCurrentCUDAStream(x.get_device()))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// Mutations through the raw multicast address ws_mc_ptr are not visible to the schema: that buffer is owned by the +// Python workspace and never aliased by the op's tensor arguments. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp4_draft_step(Tensor x, Tensor? x_rows, int rows, Tensor wf, int vocab_offset, " + "Tensor(a!) keys, " + "Tensor(b!) counters, Tensor(c!) ws_uc, int ws_mc_ptr, int rank, int world, Tensor(d!) draft_out, int step, " + "Tensor(e!) ids_next, Tensor pos_src, Tensor(f!) pos_next, Tensor(g!) kv_lens, Tensor? num_acc, " + "Tensor(h!)? seq_lens, int draft_len, int num_ctas, bool pdl, Tensor(i!)? out_keys, Tensor? exc_off=None, " + "Tensor? exc_pos=None, Tensor? exc_w=None, int exp_base=0" + ", Tensor(j!)[]? fold=None, int fold_rows=0, int fold_tokens_per_step=0, bool fixed_order=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_draft_step", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4DraftStep); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp new file mode 100644 index 000000000000..034521d29796 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp4::drafterRopeFp8Kv; + +//! The MiniMax-M3 TP4 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h): RoPE on +//! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages +//! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. +void minimaxM3Sm107Tp4DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, + torch::Tensor const& seqLens, std::optional const& specPositionOffsets, torch::Tensor& kvPages, + torch::Tensor const& blockTables, torch::Tensor& qOut, int64_t qLen, int64_t numHeadsQ, int64_t numHeadsKv, + int64_t rotaryDim, int64_t tokensPerBlock, int64_t ropeMode, bool pdl, bool maxSharedCarveout) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + 2 * numHeadsKv) * drk::kHeadDim, "qkv is not [tokens, (Hq + 2 Hkv) * 128]"); + TORCH_CHECK(qLen >= 1 && qkv.size(0) % qLen == 0, "the rows are not whole passes of q_len rows"); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= drk::kHeadDim && rotaryDim % 16 == 0, "unsupported rotary_dim"); + TORCH_CHECK(cosSin.is_cuda() && cosSin.is_contiguous() && cosSin.scalar_type() == at::kFloat + && cosSin.numel() % rotaryDim == 0, + "cos_sin must be the fp32 [positions, rotary_dim / 2] (cos, sin) table"); + int64_t const requests = qkv.size(0) / qLen; + TORCH_CHECK(seqLens.is_cuda() && seqLens.scalar_type() == at::kInt && seqLens.numel() >= requests); + if (specPositionOffsets.has_value()) + { + TORCH_CHECK(specPositionOffsets->is_cuda() && specPositionOffsets->is_contiguous() + && specPositionOffsets->scalar_type() == at::kInt && specPositionOffsets->numel() >= requests * qLen); + } + TORCH_CHECK(kvPages.is_cuda() && kvPages.is_contiguous() && kvPages.dim() == 4 + && kvPages.scalar_type() == at::kFloat8_e4m3fn && kvPages.size(1) == numHeadsKv + && kvPages.size(2) == tokensPerBlock && kvPages.size(3) == drk::kHeadDim, + "kv_pages must be the E4M3 [sub-pages, Hkv, tokens_per_block, 128] cache"); + TORCH_CHECK(blockTables.is_cuda() && blockTables.scalar_type() == at::kInt && blockTables.dim() == 3 + && blockTables.size(0) >= requests && blockTables.size(1) == 2 && blockTables.stride(2) == 1, + "block_tables must be the int32 [requests, 2, max blocks] K / V sub-page table"); + TORCH_CHECK(qOut.is_cuda() && qOut.is_contiguous() && qOut.scalar_type() == at::kFloat8_e4m3fn + && qOut.numel() == qkv.size(0) * numHeadsQ * drk::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{qkv.device()}; + drk::Params params{}; + params.qkv = reinterpret_cast<__nv_bfloat16 const*>(qkv.data_ptr()); + params.cosSin = reinterpret_cast(cosSin.data_ptr()); + params.seqLens = seqLens.data_ptr(); + params.specPositionOffsets = specPositionOffsets.has_value() ? specPositionOffsets->data_ptr() : nullptr; + params.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + params.pool = reinterpret_cast<__nv_fp8_e4m3*>(kvPages.data_ptr()); + params.blockTable = blockTables.data_ptr(); + params.blockTableRequestStride = blockTables.stride(0); + params.blockTableKvStride = blockTables.stride(1); + params.numTokens = static_cast(qkv.size(0)); + params.qLen = static_cast(qLen); + params.numHeadsQ = static_cast(numHeadsQ); + params.numHeadsKv = static_cast(numHeadsKv); + params.rotaryDim = static_cast(rotaryDim); + params.tokensPerBlock = static_cast(tokensPerBlock); + params.ropeMode = static_cast(ropeMode); + cudaError_t const err + = drk::launch(params, pdl, maxSharedCarveout, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp4_drafter_rope_fp8_kv: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp4_drafter_rope_fp8_kv(Tensor qkv, Tensor cos_sin, Tensor seq_lens, " + "Tensor? spec_position_offsets, Tensor(a!) kv_pages, Tensor block_tables, Tensor(b!) q_out, int q_len, " + "int num_heads_q, int num_heads_kv, int rotary_dim, int tokens_per_block, int rope_mode, bool pdl, bool " + "max_shared_carveout=True) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_drafter_rope_fp8_kv", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4DrafterRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp new file mode 100644 index 000000000000..3df36817814d --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp @@ -0,0 +1,89 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" +#include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The speculative step's forward-tail store and next-step staging in one launch (the C++ form of the Triton +//! kernels/m3_fold_tail.py, the same writes). The caller checks the row counts against the buffers (fold_tail). +static void foldTailImpl(bool oneWarp, torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + c10::cuda::CUDAGuard const deviceGuard{ctrl.device()}; + auto const p = minimaxM3Sm107Tp4FoldTailParams::buildParams(ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, + srcDraft, dstDraft, srcLens, dstLens, ids, draftIds, posOff, kvOff, tokensPerStep); + auto const stream = at::cuda::getCurrentCUDAStream(ctrl.get_device()); + auto const err = oneWarp + ? tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeFoldTailOneWarp(p, stream) + : tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeFoldTail(p, static_cast(rows), stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_fold_tail launch failed: ", cudaGetErrorString(err)); +} + +//! The tail as its standalone kernel (one CTA per row). +void minimaxM3Sm107Tp4FoldTail(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + foldTailImpl(false, ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, srcDraft, dstDraft, srcLens, dstLens, ids, + draftIds, posOff, kvOff, tokensPerStep); +} + +//! The tail's in-kernel form (one warp over the rows, foldTailWarp), the oracle of a caller that folds it in. +void minimaxM3Sm107Tp4FoldTailOneWarp(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + foldTailImpl(true, ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, srcDraft, dstDraft, srcLens, dstLens, ids, + draftIds, posOff, kvOff, tokensPerStep); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_fold_tail(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) dst_new, " + "Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); + m.def( + "minimax_m3_sm107_tp4_fold_tail_one_warp(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) " + "dst_new, Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_fold_tail", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4FoldTail); + m.impl("minimax_m3_sm107_tp4_fold_tail_one_warp", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4FoldTailOneWarp); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h new file mode 100644 index 000000000000..b73820b6fc54 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace minimaxM3Sm107Tp4FoldTailParams +{ + +constexpr int64_t kMaxWidth = 32; + +inline void check2d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 2 && t.size(1) <= kMaxWidth, + "minimax_m3_sm107_tp4_fold_tail: ", name, " must be a 2-D int32 CUDA view at most 32 wide"); +} + +inline void check1d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp4_fold_tail: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +//! The fold tail's operands checked and packed (the standalone fold tail op and any op that runs the tail in-kernel, +//! e.g. the drafter's last-pass draft_step, build them identically). +inline tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams buildParams(torch::Tensor const& ctrl, + torch::Tensor const& slots, int64_t rows, torch::Tensor const& srcNew, torch::Tensor& dstNew, + torch::Tensor const& srcNext, torch::Tensor& dstNext, torch::Tensor const& srcDraft, torch::Tensor& dstDraft, + torch::Tensor const& srcLens, torch::Tensor& dstLens, torch::Tensor& ids, torch::Tensor& draftIds, + torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + TORCH_CHECK(ctrl.is_cuda() && ctrl.scalar_type() == torch::kInt32 && ctrl.numel() >= 1, + "minimax_m3_sm107_tp4_fold_tail: ctrl must be an int32 CUDA tensor"); + TORCH_CHECK(slots.is_cuda() && slots.scalar_type() == torch::kInt64 && slots.dim() == 1 && slots.stride(0) == 1, + "minimax_m3_sm107_tp4_fold_tail: slots must be a contiguous 1-D int64 CUDA tensor"); + TORCH_CHECK(rows > 0 && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, + "minimax_m3_sm107_tp4_fold_tail: rows must be positive and tokens_per_step in [1, 32]"); + check2d(srcNew, "src_new"); + check2d(dstNew, "dst_new"); + check2d(srcNext, "src_next"); + check2d(dstNext, "dst_next"); + check2d(srcDraft, "src_draft"); + check2d(dstDraft, "dst_draft"); + check1d(srcLens, "src_lens"); + check1d(dstLens, "dst_lens"); + check1d(ids, "ids"); + check1d(draftIds, "draft_ids"); + check1d(posOff, "pos_off"); + check1d(kvOff, "kv_off"); + tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams p{}; + p.ctrl = ctrl.data_ptr(); + p.slots = slots.data_ptr(); + p.srcNew = srcNew.data_ptr(); + p.sn0 = srcNew.stride(0); + p.sn1 = srcNew.stride(1); + p.snW = static_cast(srcNew.size(1)); + p.dstNew = dstNew.data_ptr(); + p.dn0 = dstNew.stride(0); + p.dn1 = dstNew.stride(1); + p.dnW = static_cast(dstNew.size(1)); + p.srcNext = srcNext.data_ptr(); + p.sx0 = srcNext.stride(0); + p.sx1 = srcNext.stride(1); + p.sxW = static_cast(srcNext.size(1)); + p.dstNext = dstNext.data_ptr(); + p.dx0 = dstNext.stride(0); + p.dx1 = dstNext.stride(1); + p.dxW = static_cast(dstNext.size(1)); + p.srcDraft = srcDraft.data_ptr(); + p.sd0 = srcDraft.stride(0); + p.sd1 = srcDraft.stride(1); + p.sdW = static_cast(srcDraft.size(1)); + p.dstDraft = dstDraft.data_ptr(); + p.dd0 = dstDraft.stride(0); + p.dd1 = dstDraft.stride(1); + p.ddW = static_cast(dstDraft.size(1)); + p.srcLens = srcLens.data_ptr(); + p.dstLens = dstLens.data_ptr(); + p.ids = ids.data_ptr(); + p.draftIds = draftIds.data_ptr(); + p.posOff = posOff.data_ptr(); + p.kvOff = kvOff.data_ptr(); + p.tokensPerStep = static_cast(tokensPerStep); + return p; +} + +} // namespace minimaxM3Sm107Tp4FoldTailParams +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp new file mode 100644 index 000000000000..e2f65f63ae33 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107Tp4::fullCtxReduce; + +//! Merges the PrimTS split-KV decode scratch (partial_o [B, Hkv, S, R, 128] 16-bit, partial_stats [B, Hkv, S, R] +//! fp32 log2-LSE) of the MiniMax-M3 TP4 full-context decode into out (BF16 [B * SQ, Hkv * hR, 128]); with q8 / sf +//! also its MXFP8 copy (E4M3 [>= B * SQ, Hkv * hR * 128] + 128x4-swizzled UE8M0 scales, as mxfp8_quantize(out, True)). +void minimaxM3Sm107Tp4FullCtxReduce(torch::Tensor const& partialO, torch::Tensor const& partialStats, + torch::Tensor const& seqLensKv, torch::Tensor& out, int64_t seqLenQ, int64_t qTokensPerCta, int64_t tileKv, + int64_t instsKv, bool causal, bool pdl, std::optional const& q8, + std::optional const& sf) +{ + TORCH_CHECK(partialO.is_cuda() && partialO.is_contiguous() && partialO.dim() == 5); + TORCH_CHECK(partialO.scalar_type() == at::kHalf || partialO.scalar_type() == at::kBFloat16); + TORCH_CHECK(partialO.size(4) == fcr::kHeadDim && partialO.size(2) <= fcr::kMaxSplits); + TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() + && partialStats.sizes() == partialO.sizes().slice(0, 4)); + TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + int const rows = static_cast(partialO.size(3)); + TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); + int const headsPerKv = rows / static_cast(seqLenQ); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.is_contiguous() + && out.numel() == partialO.size(0) * seqLenQ * partialO.size(1) * headsPerKv * fcr::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{partialO.device()}; + fcr::Params params{partialO.data_ptr(), partialStats.data_ptr(), seqLensKv.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), static_cast(partialO.size(0)), + static_cast(partialO.size(1)), static_cast(partialO.size(2)), rows, static_cast(seqLenQ), + headsPerKv, static_cast(qTokensPerCta), static_cast(tileKv), static_cast(instsKv), causal, + partialO.scalar_type() == at::kBFloat16}; + TORCH_CHECK(q8.has_value() == sf.has_value()); + if (q8.has_value()) + { + int64_t const width = partialO.size(1) * headsPerKv * fcr::kHeadDim; + int64_t const tokens = partialO.size(0) * seqLenQ; + int64_t const tileRows = (tokens + 127) / 128 * 128; + TORCH_CHECK(q8->is_cuda() && q8->is_contiguous() && q8->scalar_type() == at::kFloat8_e4m3fn && q8->dim() == 2 + && q8->size(1) == width && q8->size(0) >= tokens); + TORCH_CHECK(sf->is_cuda() && sf->is_contiguous() && sf->scalar_type() == at::kByte + && sf->numel() >= tileRows * (width / 32)); + params.q8 = reinterpret_cast(q8->data_ptr()); + params.sf = sf->data_ptr(); + params.numSfKTiles = static_cast(width / 128); + } + cudaError_t const err + = fcr::invokeFullCtxReduce(params, pdl, at::cuda::getCurrentCUDAStream(partialO.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp4_fullctx_reduce: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp4_fullctx_reduce(Tensor partial_o, Tensor partial_stats, Tensor seq_lens_kv, " + "Tensor(a!) out, int seq_len_q, int q_tokens_per_cta, int tile_kv, int insts_kv, bool causal, bool pdl, " + "Tensor(b!)? q8=None, Tensor(c!)? sf=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_fullctx_reduce", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4FullCtxReduce); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4McAllreduceNormOp.cpp new file mode 100644 index 000000000000..27442555f3b4 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4McAllreduceNormOp.cpp @@ -0,0 +1,181 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h" + +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxTokens = 48; +constexpr int64_t kCluster = 8; +constexpr int64_t kMaxWorld = 8; + +} // namespace + +//! MiniMax-M3 SM107 TP4 multicast all-reduce (stock rounding points: BF16 rank sum, then the residual add) + residual + +//! RMSNorm (+ optional MoE finalize, + optional MXFP8 output). +//! +//! @param x BF16 [M, H] partial, or with `perm` the permuted FC2 output [rows, H]. +//! @param x2 Optional BF16 [M, H] second addend (the shared expert in finalize mode). +//! @param residual BF16 [M, H], 1 <= M <= 48. +//! @param normWeight BF16 [H]. +//! @param wsUc This rank's unicast view of the exchange workspace (int32 words); re-armed in place. +//! @param wsMcPtr Device address of the workspace's multicast mapping. +//! @param flags Int32 per-CTA call-parity counters, at least maxt * 8; updated in place. +//! @param cluster Cluster size; 8 is the only one built. +//! @param perm, scales Finalize mode: int32 [M, topk] permuted-row index and fp32 / bf16 [M, topk] routing weights. +//! @param maxt Token capacity of the workspace bucket (8, 16 or 48). +//! @param quant Also return the MXFP8 copy of the normed rows (E4M3 [M, H], swizzled UE8M0 scales). +//! @param prepushed 0, or 2 / 3 when producers already pushed this rank's rows (see the kernel header). +//! @param opts Epilogue variant; every variant is bit-identical. +//! @param capOut Optional BF16 [>= M, H] view (row stride a multiple of 8, 16-byte aligned) that receives the Eagle3 +//! capture of the updated residual: m3_spec_tail capture_norm with capW [H] and capEps, bit for bit (H 6144; not the +//! drafter boundary, opts bit 16). +//! @return [normed, updated residual] or, with `quant`, [normed, updated residual, E4M3 rows, scale bytes]. +std::vector minimaxM3Sm107Tp4McAllreduceNorm(torch::Tensor const& x, + std::optional const& x2, torch::Tensor const& residual, torch::Tensor const& normWeight, double eps, + torch::Tensor& wsUc, int64_t wsMcPtr, torch::Tensor& flags, int64_t rank, int64_t world, int64_t cluster, bool pdl, + std::optional const& perm, std::optional const& scales, int64_t maxt, bool quant, + int64_t prepushed, int64_t opts, std::optional const& capOut, + std::optional const& capW, double capEps) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 107, "minimax_m3_sm107_tp4_mc_allreduce_norm requires SM107, got SM", sm); + TORCH_CHECK(cluster == kCluster, "minimax_m3_sm107_tp4_mc_allreduce_norm: only cluster 8 is built, got ", cluster); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107Tp4McAllreduceNormSupported( + static_cast(maxt), static_cast(opts)), + "minimax_m3_sm107_tp4_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); + TORCH_CHECK( + world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); + + tensorrt_llm::kernels::MinimaxM3Sm107Tp4McAllreduceNormParams p; + if (x2.has_value()) + { + TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), + "x2 must be a contiguous bf16 tensor shaped like residual"); + p.x2 = reinterpret_cast<__nv_bfloat16 const*>(x2->data_ptr()); + } + if (perm.has_value()) + { + // Finalize mode: x is the permuted FC2 output [rows, H]. + TORCH_CHECK(scales.has_value() && perm->scalar_type() == at::kInt && perm->is_contiguous() && perm->dim() == 2 + && perm->size(0) == residual.size(0), + "perm must be a contiguous int32 [M, topk] tensor, with scales"); + TORCH_CHECK(scales->is_contiguous() && scales->numel() == perm->numel() + && (scales->scalar_type() == at::kFloat || scales->scalar_type() == at::kBFloat16), + "scales must be a contiguous fp32 or bf16 tensor with one weight per perm entry"); + TORCH_CHECK(x.dim() == 2 && x.size(1) == residual.size(1), "fc2 must be [rows, H]"); + p.fc2 = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.perm = perm->data_ptr(); + p.scales = scales->data_ptr(); + p.scaleFp32 = scales->scalar_type() == at::kFloat; + p.topk = static_cast(perm->size(1)); + } + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2, + "x must be a contiguous 2-D bf16 CUDA tensor"); + TORCH_CHECK(residual.dim() == 2 && residual.size(0) >= 1 && residual.size(0) <= std::min(kMaxTokens, maxt), + "residual must hold 1..min(48, bucket) rows"); + TORCH_CHECK(residual.scalar_type() == at::kBFloat16 && residual.is_contiguous() + && (perm.has_value() || residual.sizes() == x.sizes()), + "residual must be a contiguous bf16 tensor shaped like x"); + TORCH_CHECK( + normWeight.scalar_type() == at::kBFloat16 && normWeight.numel() == residual.size(1), "norm_w must be bf16 [H]"); + TORCH_CHECK(wsUc.scalar_type() == at::kInt && wsUc.is_cuda(), "ws_uc must be an int32 CUDA tensor"); + TORCH_CHECK(flags.scalar_type() == at::kInt && flags.is_contiguous() && flags.numel() >= maxt * kCluster, + "flags must be a contiguous int32 tensor of at least bucket * 8 counters"); + + int64_t const m = residual.size(0), h = residual.size(1); + TORCH_CHECK(!quant || h % 128 == 0, "MXFP8 output needs H % 128 == 0"); + c10::cuda::CUDAGuard const deviceGuard{residual.device()}; + auto outNormed = torch::empty_like(residual); + auto outResidual = torch::empty_like(residual); + auto outQ = quant ? torch::empty({m, h}, residual.options().dtype(at::kFloat8_e4m3fn)) : torch::Tensor(); + auto outSf = quant ? torch::empty({(m + 127) / 128 * 128 * (h / 32)}, residual.options().dtype(at::kByte)) + : torch::Tensor(); + + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.numTokens = static_cast(m); + p.hidden = static_cast(h); + p.residual = reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()); + p.normWeight = reinterpret_cast<__nv_bfloat16 const*>(normWeight.data_ptr()); + p.eps = static_cast(eps); + p.wsUc = wsUc.data_ptr(); + p.wsMc = reinterpret_cast(wsMcPtr); + p.flags = flags.data_ptr(); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.outNormed = reinterpret_cast<__nv_bfloat16*>(outNormed.data_ptr()); + p.outResidual = reinterpret_cast<__nv_bfloat16*>(outResidual.data_ptr()); + p.outQuant = quant ? reinterpret_cast(outQ.data_ptr()) : nullptr; + p.outSf = quant ? outSf.data_ptr() : nullptr; + p.prepushed = static_cast(prepushed); + if (capOut.has_value()) + { + TORCH_CHECK((opts & 16) == 0, "the capture epilogue is not built for the drafter boundary (opts bit 16)"); + TORCH_CHECK(capOut->scalar_type() == at::kBFloat16 && capOut->dim() == 2 && capOut->size(0) >= m + && capOut->size(1) == h && capOut->stride(1) == 1 && capOut->stride(0) % 8 == 0 + && reinterpret_cast(capOut->data_ptr()) % 16 == 0 && capOut->device() == residual.device(), + "cap_out must be a bf16 [>= M, H] view with unit column stride, a row stride that is a multiple of 8 and " + "16-byte alignment"); + TORCH_CHECK( + capW.has_value() && capW->scalar_type() == at::kBFloat16 && capW->is_contiguous() && capW->numel() == h, + "cap_w must be a contiguous bf16 [H] tensor"); + p.capOut = reinterpret_cast<__nv_bfloat16*>(capOut->data_ptr()); + p.capStride = capOut->stride(0); + p.capWeight = reinterpret_cast<__nv_bfloat16 const*>(capW->data_ptr()); + p.capEps = static_cast(capEps); + } + tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp4McAllreduceNorm(p, static_cast(maxt), + static_cast(opts), pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); + if (quant) + { + return {outNormed, outResidual, outQ, outSf}; + } + return {outNormed, outResidual}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_mc_allreduce_norm(Tensor x, Tensor? x2, Tensor residual, Tensor norm_w, float eps, " + "Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) flags, int rank, int world, int cluster, bool pdl, " + "Tensor? perm, Tensor? scales, int maxt, bool quant, int prepushed, int opts, Tensor(c!)? cap_out=None, " + "Tensor? cap_w=None, float cap_eps=0.) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_mc_allreduce_norm", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4McAllreduceNorm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp new file mode 100644 index 000000000000..3ae68e7b4681 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp @@ -0,0 +1,185 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp4::producerPw4; + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPageSize = 128; + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "sm107_tp4_producer_pw: ", name, " must be a CUDA tensor"); +} + +} // namespace + +//! The SM107 TP4 NVFP4 QKV + indexer norm/RoPE/KV-insert producer (same arguments and outputs as +//! minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert at 16 query, 1 KV and 1 index head, head_dim 128, rotary_dim 64): +//! returns (q FP8 [T, 8, 128], index_q FP8 [T, 1, 128]) and writes the NVFP4 K/V of the KV head and the FP8 index-K +//! of each token's cache slot. `early` triggers the dependents on entry; `prewait_step_inputs` reads the position ids, +//! the slots and the RoPE rows before the grid-dependency wait (none of them may be written by a PDL predecessor). +std::vector minimaxM3Sm107Tp4ProducerPw(torch::Tensor const& packed, torch::Tensor& kvData, + torch::Tensor& kvScale, torch::Tensor& indexK, torch::Tensor const& outCacheLoc, torch::Tensor const& kvQuantScale, + double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds, bool pdl, + bool early, bool prewaitStepInputs) +{ + checkCuda(packed, "packed"); + for (torch::Tensor const* t : std::initializer_list{&kvData, &kvScale, &indexK, &outCacheLoc, + &kvQuantScale, &qWeight, &kWeight, &indexQWeight, &indexKWeight, &rotaryCosSin, &positionIds}) + { + TORCH_CHECK(t->is_cuda() && t->get_device() == packed.get_device(), + "sm107_tp4_producer_pw: every operand must be on packed's CUDA device"); + } + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0 + && reinterpret_cast(kvData.data_ptr()) % 2 == 0 + && reinterpret_cast(indexK.data_ptr()) % 4 == 0, + "sm107_tp4_producer_pw: packed and the caches must start at the vector-load / store alignment"); + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp4_producer_pw requires SM100 or newer"); + int64_t const headsQ = pw::kHeadsQ; + int64_t const headsKV = pw::kHeadsKV; + int64_t const headsIndex = pw::kHeadsIndex; + int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; + TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim + && packed.is_contiguous(), + "sm107_tp4_producer_pw: packed must be contiguous BF16 [T, 20 * 128]"); + // The kernel's layout (the native producer's checkMinimaxM3Nvfp4HndKVPools at one KV head): HND pools + // [pages, 2 planes, 1 head, 128 tokens, 64 data / 8 scale bytes], dense token rows and head blocks, planes and + // pages that do not overlap. + TORCH_CHECK(kvData.scalar_type() == at::kByte && kvData.dim() == 5 && kvData.size(0) > 0 && kvData.size(1) == 2 + && kvData.size(2) == headsKV && kvData.size(3) == kPageSize && kvData.size(4) == kHeadDim / 2, + "sm107_tp4_producer_pw: kv_data must be the NVFP4 pool [pages, 2, 1, 128, 64] uint8"); + TORCH_CHECK(kvScale.scalar_type() == at::kByte && kvScale.dim() == 5 && kvScale.size(0) == kvData.size(0) + && kvScale.size(1) == 2 && kvScale.size(2) == headsKV && kvScale.size(3) == kPageSize + && kvScale.size(4) == kHeadDim / 16, + "sm107_tp4_producer_pw: kv_scale must be the E4M3 scale pool [pages, 2, 1, 128, 8] uint8"); + TORCH_CHECK(kvData.stride(4) == 1 && kvData.stride(3) == kvData.size(4) + && kvData.stride(2) == kvData.size(3) * kvData.stride(3), + "sm107_tp4_producer_pw: kv_data must have dense token rows and head blocks"); + TORCH_CHECK(kvScale.stride(4) == 1 && kvScale.stride(3) == kvScale.size(4) + && kvScale.stride(2) == kvScale.size(3) * kvScale.stride(3), + "sm107_tp4_producer_pw: kv_scale must have dense token rows and head blocks"); + TORCH_CHECK(kvData.stride(1) >= kvData.size(2) * kvData.stride(2) && kvData.stride(0) >= 2 * kvData.stride(1) + && kvScale.stride(1) >= kvScale.size(2) * kvScale.stride(2) && kvScale.stride(0) >= 2 * kvScale.stride(1), + "sm107_tp4_producer_pw: pool planes or pages overlap"); + TORCH_CHECK(kvData.stride(0) % 2 == 0 && kvData.stride(1) % 2 == 0, + "sm107_tp4_producer_pw: kv_data page and plane strides must keep the uint16 stores aligned"); + TORCH_CHECK(indexK.scalar_type() == at::kFloat8_e4m3fn && indexK.dim() == 4 && indexK.size(0) == kvData.size(0) + && indexK.size(1) == 1 && indexK.size(2) == kPageSize && indexK.size(3) == kHeadDim && indexK.stride(3) == 1 + && indexK.stride(2) == kHeadDim && indexK.stride(0) >= kPageSize * kHeadDim && indexK.stride(0) % 4 == 0, + "sm107_tp4_producer_pw: index_k must be FP8 [pages, 1, 128, 128] with a page stride keeping the 32-bit stores " + "aligned"); + int64_t const numTokens = packed.size(0); + TORCH_CHECK( + outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous() && outCacheLoc.numel() >= numTokens, + "sm107_tp4_producer_pw: out_cache_loc must be int32 [>= T]"); + TORCH_CHECK( + positionIds.scalar_type() == at::kInt && positionIds.is_contiguous() && positionIds.numel() >= numTokens, + "sm107_tp4_producer_pw: position_ids must be int32 [>= T]"); + TORCH_CHECK(kvQuantScale.scalar_type() == at::kFloat && kvQuantScale.is_contiguous() && kvQuantScale.numel() >= 3, + "sm107_tp4_producer_pw: kv_quant_scale must be fp32 [>= 3]"); + TORCH_CHECK(rotaryCosSin.scalar_type() == at::kFloat && rotaryCosSin.is_contiguous() && rotaryCosSin.dim() == 3 + && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == 32 + && reinterpret_cast(rotaryCosSin.data_ptr()) % 16 == 0, + "sm107_tp4_producer_pw: rotary_cos_sin must be fp32 [positions, 2, 32], 16-byte aligned (read as float4)"); + for (auto const* w : {&qWeight, &kWeight, &indexQWeight, &indexKWeight}) + { + TORCH_CHECK(w->scalar_type() == at::kBFloat16 && w->is_contiguous() && w->numel() == kHeadDim + && reinterpret_cast(w->data_ptr()) % 8 == 0, + "sm107_tp4_producer_pw: norm weights must be BF16 [128], 8-byte aligned (read as 4-value words)"); + } + auto const fp8 = packed.options().dtype(at::kFloat8_e4m3fn); + auto qOut = torch::empty({numTokens, headsQ, kHeadDim}, fp8); + auto indexQOut = torch::empty({numTokens, headsIndex, kHeadDim}, fp8); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + pw::Params p{}; + p.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + p.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + p.indexQOut = reinterpret_cast<__nv_fp8_e4m3*>(indexQOut.data_ptr()); + p.kvData = kvData.data_ptr(); + p.kvScale = kvScale.data_ptr(); + p.indexK = reinterpret_cast<__nv_fp8_e4m3*>(indexK.data_ptr()); + p.outCacheLoc = outCacheLoc.data_ptr(); + p.kvQuantScale = kvQuantScale.data_ptr(); + p.dataPageStride = kvData.stride(0); + p.dataPlaneStride = kvData.stride(1); + p.dataHeadStride = kvData.stride(2); + p.dataTokenStride = kvData.stride(3); + p.scalePageStride = kvScale.stride(0); + p.scalePlaneStride = kvScale.stride(1); + p.scaleHeadStride = kvScale.stride(2); + p.indexPageStride = indexK.stride(0); + p.indexTokenStride = indexK.stride(2); + p.numPages = kvData.size(0); + p.numTokens = static_cast(numTokens); + p.eps = static_cast(eps); + p.qWeight = reinterpret_cast<__nv_bfloat16 const*>(qWeight.data_ptr()); + p.kWeight = reinterpret_cast<__nv_bfloat16 const*>(kWeight.data_ptr()); + p.indexQWeight = reinterpret_cast<__nv_bfloat16 const*>(indexQWeight.data_ptr()); + p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); + p.rotaryCosSin = rotaryCosSin.data_ptr(); + p.positionIds = positionIds.data_ptr(); + cudaError_t const err = pw::invokeProducerPw4(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "sm107_tp4_producer_pw launch: ", cudaGetErrorString(err)); + } + return {qOut, indexQOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_producer_pw(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) index_k, " + "Tensor out_cache_loc, Tensor kv_quant_scale, float eps, Tensor q_weight, Tensor k_weight, " + "Tensor index_q_weight, Tensor index_k_weight, Tensor rotary_cos_sin, Tensor position_ids, bool pdl, " + "bool early, bool prewait_step_inputs=False) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_producer_pw", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4ProducerPw); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp new file mode 100644 index 000000000000..4f7fda5cac90 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp @@ -0,0 +1,113 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3Sm107Tp4::sampledAccept; + +//! uint64 words of the sampled-accept exchange workspace for `world` ranks and SA-1R variant `oneRoundK`. +int64_t minimaxM3Sm107Tp4SampledAcceptWsWords(int64_t world, int64_t oneRoundK) +{ + return sa::sampledAcceptWsWords(static_cast(world), static_cast(oneRoundK)); +} + +//! Sampled (temperature + top-p) verification from this rank's BF16 logits shard [rows, Vs], rows = nreq * +//! (draft_len + 1) <= 48: the sampled token of every row (tokens), then the strict acceptance against draft (accTokens, +//! numAcc, gatherIds, optional nextNew column 0). Per-row Philox (seeds, offsets); the exchange goes through the +//! multicast workspace wsMcPtr / wsUc with the (finished rows, call) counters. +void minimaxM3Sm107Tp4SampledAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor const& temperatures, + torch::Tensor const& topPs, torch::Tensor const& seeds, torch::Tensor const& offsets, torch::Tensor& wsUc, + int64_t wsMcPtr, torch::Tensor& counters, torch::Tensor& tokens, int64_t rank, int64_t world, + torch::Tensor const& draft, torch::Tensor const& batchIdx, torch::Tensor& accTokens, torch::Tensor& numAcc, + torch::Tensor& gatherIds, int64_t forceTotal, std::optional const& nextNew, + std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, + bool prewaitParams, bool earlyTrigger, int64_t oneRoundK, std::optional const& decided, bool ilpWalk) +{ + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + sa::Args a; + a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + a.ld = logits.stride(0); + a.rows = static_cast(logits.size(0)); + a.Vs = static_cast(logits.size(1)); + TORCH_CHECK(a.Vs <= sa::kMaxVs && a.rows <= sa::kMaxRows && world <= sa::kMaxWorld + && static_cast(vocabOffset) + a.Vs <= (1LL << sa::kIdBits)); + a.vocab_offset = static_cast(vocabOffset); + a.temperatures = temperatures.data_ptr(); + a.top_ps = topPs.data_ptr(); + a.seeds = reinterpret_cast(seeds.data_ptr()); + a.offsets = reinterpret_cast(offsets.data_ptr()); + a.ws_uc = reinterpret_cast(wsUc.data_ptr()); + a.ws_mc = reinterpret_cast(wsMcPtr); + a.counters = reinterpret_cast(counters.data_ptr()); + a.tokens = tokens.data_ptr(); + a.extra_attempts = extraAttempts.has_value() ? extraAttempts->data_ptr() : nullptr; + a.rank = static_cast(rank); + a.world = static_cast(world); + a.draft = draft.data_ptr(); + a.batch_idx = batchIdx.data_ptr(); + a.acc_tokens = accTokens.data_ptr(); + a.num_acc = numAcc.data_ptr(); + a.gather_ids = reinterpret_cast(gatherIds.data_ptr()); + a.draft_len = static_cast(draft.size(1)); + a.nreq = static_cast(draft.size(0)); + a.force_total = static_cast(forceTotal); + a.next_new = nextNew.has_value() ? nextNew->data_ptr() : nullptr; + a.next_stride = nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0; + a.force_vec = forceVec.has_value() ? forceVec->data_ptr() : nullptr; + a.prewait_params = prewaitParams ? 1 : 0; + a.early_trigger = earlyTrigger ? 1 : 0; + TORCH_CHECK(oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "one_round_k must be 0, 9 or 82, got ", oneRoundK); + a.one_round_k = static_cast(oneRoundK); + a.decided = decided.has_value() ? decided->data_ptr() : nullptr; + a.ilp_walk = ilpWalk ? 1 : 0; + cudaError_t const err = sa::invokeSampledAccept(a, usePdl, at::cuda::getCurrentCUDAStream(logits.get_device())); + TORCH_CHECK(err == cudaSuccess, "m3_sampled_accept launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// ws_mc_ptr is the multicast address of ws_uc's mapping: the op writes through it, which the schema cannot show. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + // No tensor arguments, so no dispatch key: registered with the schema as a catch-all kernel. + m.def("minimax_m3_sm107tp4_sampled_accept_ws_words(int world, int one_round_k=0) -> int", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SampledAcceptWsWords); + m.def( + "minimax_m3_sm107tp4_sampled_accept(Tensor logits, int vocab_offset, Tensor temperatures, Tensor top_ps, " + "Tensor seeds, Tensor offsets, Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) counters, Tensor(c!) tokens, " + "int rank, int world, Tensor draft, Tensor batch_idx, Tensor(d!) acc_tokens, Tensor(e!) num_acc, " + "Tensor(f!) gather_ids, int force_total, Tensor(g!)? next_new, Tensor? force_vec, " + "Tensor(h!)? extra_attempts, bool use_pdl, bool prewait_params=True, bool early_trigger=False, " + "int one_round_k=0, Tensor(i!)? decided=None, bool ilp_walk=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_sampled_accept", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SampledAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SelectLongOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SelectLongOp.cpp new file mode 100644 index 000000000000..b1cf6be45569 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SelectLongOp.cpp @@ -0,0 +1,87 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, into a caller-owned table. +//! +//! @param scores Float32 CUDA tensor [num_kv_heads, num_blocks, total_q], any strides. +//! @param nValidBlocks Contiguous int32 CUDA tensor with at least total_q elements. +//! @param out Int32 CUDA tensor of logical shape [total_q, num_kv_heads, 16] with unit stride in the last dim (the +//! head-major backing the target hands over is a permuted view). Every entry is written. +//! @param triggerPhase Where the dependent grid is released without earlyTrigger: 0 at this grid's exit, 1 right +//! after every score of the row is loaded (both after this grid's dependency wait). +//! @param fastTail Rank rows of at most 64 candidates on warp 0 and read the valid count before the wait (same picks). +void minimaxM3Sm107Tp4SelectLong(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, torch::Tensor& out, + int64_t initBlocks, int64_t localBlocks, bool usePdl, bool earlyTrigger, int64_t triggerPhase, bool fastTail) +{ + TORCH_CHECK(scores.is_cuda() && scores.scalar_type() == torch::kFloat32 && scores.dim() == 3, + "minimax_m3_sm107_tp4_select_long: scores must be a 3-D float32 CUDA tensor [num_kv_heads, num_blocks, " + "total_q]"); + TORCH_CHECK(nValidBlocks.is_cuda() && nValidBlocks.scalar_type() == torch::kInt32 && nValidBlocks.is_contiguous() + && nValidBlocks.numel() >= scores.size(2), + "minimax_m3_sm107_tp4_select_long: n_valid_blocks must be contiguous int32 CUDA with at least total_q " + "elements"); + int64_t const numKvHeads = scores.size(0); + int64_t const numBlocks = scores.size(1); + int64_t const totalQ = scores.size(2); + TORCH_CHECK(numBlocks <= tensorrt_llm::kernels::minimaxM3Sm107Tp4::kSelectLong2MaxBlocks, + "minimax_m3_sm107_tp4_select_long supports at most ", + tensorrt_llm::kernels::minimaxM3Sm107Tp4::kSelectLong2MaxBlocks, " blocks, got ", numBlocks); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kInt32 && out.dim() == 3 && out.size(0) == totalQ + && out.size(1) == numKvHeads && out.size(2) == 16 && out.stride(2) == 1, + "minimax_m3_sm107_tp4_select_long: out must be int32 [total_q, num_kv_heads, 16] with unit last-dim stride"); + + c10::cuda::CUDAGuard const deviceGuard{scores.device()}; + auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device()); + cudaError_t const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSelectLong2(scores.data_ptr(), + scores.stride(0), scores.stride(1), scores.stride(2), nValidBlocks.data_ptr(), out.data_ptr(), + out.stride(0), out.stride(1), static_cast(numKvHeads), static_cast(numBlocks), + static_cast(totalQ), static_cast(initBlocks), static_cast(localBlocks), usePdl, + earlyTrigger, static_cast(triggerPhase), fastTail, stream); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_sm107_tp4_select_long launch failed: ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_select_long(Tensor scores, Tensor n_valid_blocks, Tensor(a!) out, int init_blocks, " + "int local_blocks, bool use_pdl=False, bool early_trigger=False, int trigger_phase=0, bool fast_tail=True) -> " + "()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_select_long", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SelectLong); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp new file mode 100644 index 000000000000..b312cf79a715 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp4/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, +//! d]. +void minimaxM3Sm107Tp4SiluAndMul(torch::Tensor const& x, torch::Tensor& out) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp4_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp4_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp4_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); + TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp4_silu_and_mul: at most 65535 rows"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err + = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), + reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_silu_and_mul launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SiluAndMul); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp new file mode 100644 index 000000000000..53121e3b1e7a --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp @@ -0,0 +1,191 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeads = 16; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kTopK = 16; + +void checkBytes(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, "minimax_m3_sm107_tp4_sparse_decode_lc: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3 long-context sparse decode over the selected NVFP4 blocks (TP4 geometry: 16 query heads on +//! one KV head, head_dim 128, page 128, top-k 16). +//! +//! @param q E4M3 bytes [T, 16, 128], unit stride in the last dim. +//! @param k, v Packed E2M1 bytes [pages, 1, 128, 64] (the per-layer HND views). +//! @param kBlockScale, vBlockScale E4M3 bytes [pages, 1, 128, 8]: K token-major, V in 4x4 token-quad order. +//! @param kGlobalScale, vGlobalScale One-element float32 per-tensor dequant scales. +//! @param topk Int32 [T, 16] ascending selected block ids with -1 tail (kv head 0 of the head-major table). +//! @param out BF16 [T, 2048], written. q8 / sf, when given, receive the MXFP8 copy and its swizzled UE8M0 scales. +void minimaxM3Sm107Tp4SparseDecodeLc(torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& kBlockScale, torch::Tensor const& vBlockScale, torch::Tensor const& kGlobalScale, + torch::Tensor const& vGlobalScale, torch::Tensor const& topk, torch::Tensor const& blockTable, + torch::Tensor const& seqLens, torch::Tensor& out, std::optional q8, std::optional sf, + int64_t numSfKTiles, int64_t decodeQueryLen, double smScale, bool linearSoftmax, bool usePdl, bool earlyTrigger, + std::optional prefetchA, std::optional prefetchB, std::optional n1Self, + std::optional n1Next, std::optional n1K, std::optional n1V, + std::optional n1KScale, std::optional n1VScale, bool identity, bool oneDiv, + bool identityStep, bool btPrewait) +{ + TORCH_CHECK(earlyTrigger, "minimax_m3_sm107_tp4_sparse_decode_lc is instantiated with early_trigger=True only"); + TORCH_CHECK(q.is_cuda(), "minimax_m3_sm107_tp4_sparse_decode_lc: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m3_sm107_tp4_sparse_decode_lc requires an SM100-family GPU"); + checkBytes(q, "q"); + checkBytes(k, "k"); + checkBytes(v, "v"); + checkBytes(kBlockScale, "k_block_scale"); + checkBytes(vBlockScale, "v_block_scale"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeads && q.size(2) == kHeadDim && q.stride(2) == 1, + "minimax_m3_sm107_tp4_sparse_decode_lc: q must be [T, 16, 128] with unit last-dim stride"); + TORCH_CHECK(kGlobalScale.scalar_type() == torch::kFloat32 && vGlobalScale.scalar_type() == torch::kFloat32, + "minimax_m3_sm107_tp4_sparse_decode_lc: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m3_sm107_tp4_sparse_decode_lc: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_sm107_tp4_sparse_decode_lc: block_table and seq_lens must be int32"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.size(0) == q.size(0) + && out.size(1) == kHeads * kHeadDim && out.stride(1) == 1, + "minimax_m3_sm107_tp4_sparse_decode_lc: out must be bfloat16 [T, 2048] with unit last-dim stride"); + TORCH_CHECK( + q8.has_value() == sf.has_value(), "minimax_m3_sm107_tp4_sparse_decode_lc: pass both q8 and sf, or neither"); + bool const mxfp8 = q8.has_value(); + if (mxfp8) + { + checkBytes(*q8, "q8"); + checkBytes(*sf, "sf"); + } + + tensorrt_llm::kernels::minimaxM3Sm107Tp4::MinimaxM3SparseDecodeLcParams prm{}; + prm.q = reinterpret_cast(q.data_ptr()); + prm.q_tok = q.stride(0); + prm.q_head = q.stride(1); + prm.k = reinterpret_cast(k.data_ptr()); + prm.v = reinterpret_cast(v.data_ptr()); + prm.k_blk = k.stride(0); + prm.v_blk = v.stride(0); + prm.ksc = reinterpret_cast(kBlockScale.data_ptr()); + prm.vsc = reinterpret_cast(vBlockScale.data_ptr()); + prm.ksc_blk = kBlockScale.stride(0); + prm.vsc_blk = vBlockScale.stride(0); + prm.kgs = kGlobalScale.data_ptr(); + prm.vgs = vGlobalScale.data_ptr(); + prm.topk = topk.data_ptr(); + prm.t_tok = topk.stride(0); + prm.t_k = topk.stride(1); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.seq_lens = seqLens.data_ptr(); + prm.dql = static_cast(decodeQueryLen); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + prm.out_stride = out.stride(0); + prm.q8 = mxfp8 ? reinterpret_cast(q8->data_ptr()) : nullptr; + prm.sf = mxfp8 ? reinterpret_cast(sf->data_ptr()) : nullptr; + prm.num_sf_k_tiles = static_cast(numSfKTiles); + prm.identity = identity ? 1 : 0; + prm.identity_step = (identity || identityStep) ? 1 : 0; + prm.one_div = oneDiv ? 1 : 0; + prm.bt_prewait = btPrewait ? 1 : 0; + auto const prefetchRange = [](std::optional const& t, uint8_t const*& ptr, int64_t& bytes) + { + ptr = nullptr; + bytes = 0; + if (t.has_value()) + { + TORCH_CHECK(t->is_cuda() && t->is_contiguous(), + "minimax_m3_sm107_tp4_sparse_decode_lc: a prefetch range must be a contiguous CUDA tensor"); + ptr = reinterpret_cast(t->data_ptr()); + bytes = t->numel() * static_cast(t->element_size()); + } + }; + prefetchRange(prefetchA, prm.pf0, prm.pf0_bytes); + prefetchRange(prefetchB, prm.pf1, prm.pf1_bytes); + bool const n1 = n1Self.has_value(); + TORCH_CHECK(n1 == n1Next.has_value() && n1 == n1K.has_value() && n1 == n1V.has_value() && n1 == n1KScale.has_value() + && n1 == n1VScale.has_value(), + "minimax_m3_sm107_tp4_sparse_decode_lc: pass all six next-layer prefetch tensors, or none"); + if (n1) + { + for (auto const* t : {&*n1Self, &*n1Next}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == torch::kInt32 && t->is_contiguous() && t->dim() == 2 + && t->size(1) == kTopK && t->size(0) >= q.size(0), + "minimax_m3_sm107_tp4_sparse_decode_lc: the id buffers must be contiguous int32 [>= T, 16]"); + } + TORCH_CHECK(n1K->stride(0) == k.stride(0) && n1V->stride(0) == v.stride(0) + && n1KScale->stride(0) == kBlockScale.stride(0) && n1VScale->stride(0) == vBlockScale.stride(0) + && n1K->size(0) == k.size(0), + "minimax_m3_sm107_tp4_sparse_decode_lc: the next layer's pools must have this layer's page strides"); + prm.n1_self = n1Self->data_ptr(); + prm.n1_next = n1Next->data_ptr(); + prm.n1_k = reinterpret_cast(n1K->data_ptr()); + prm.n1_v = reinterpret_cast(n1V->data_ptr()); + prm.n1_ksc = reinterpret_cast(n1KScale->data_ptr()); + prm.n1_vsc = reinterpret_cast(n1VScale->data_ptr()); + prm.n1_pages = n1K->size(0); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSparseDecodeLc4( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, earlyTrigger, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_sparse_decode_lc launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_sparse_decode_lc(Tensor q, Tensor k, Tensor v, Tensor k_block_scale, Tensor " + "v_block_scale, " + "Tensor k_global_scale, Tensor v_global_scale, Tensor topk, Tensor block_table, Tensor seq_lens, " + "Tensor(a!) out, Tensor(b!)? q8, Tensor(c!)? sf, int num_sf_k_tiles, int decode_query_len, float sm_scale, " + "bool linear_softmax, bool use_pdl, bool early_trigger, Tensor? prefetch_a=None, Tensor? prefetch_b=None, " + "Tensor(d!)? n1_self=None, Tensor? n1_next=None, Tensor? n1_k=None, Tensor? n1_v=None, Tensor? " + "n1_k_scale=None, " + "Tensor? n1_v_scale=None, bool identity=False, bool one_div=False, bool identity_step=False, " + "bool bt_prewait=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_sparse_decode_lc", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SparseDecodeLc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp new file mode 100644 index 000000000000..2142baa18599 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void check(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp4_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +} // namespace + +//! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch +//! (kernels/minimaxM3Sm107Tp4/specOffsets.h). +void minimaxM3Sm107Tp4SpecOffsets( + torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) +{ + check(pos, "pos"); + check(posOff, "pos_off"); + check(kv, "kv"); + check(kvOff, "kv_off"); + TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), + "minimax_m3_sm107_tp4_spec_offsets: each tensor and its offsets must have the same length"); + if (pos.numel() == 0 && kv.numel() == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{pos.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSpecOffsets(pos.data_ptr(), + posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), + at::cuda::getCurrentCUDAStream(pos.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_spec_offsets launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SpecOffsets); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp new file mode 100644 index 000000000000..59009e461d33 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp4/stepPrologue.h). +//! +//! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the +//! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses +//! the caller guarantees are live. +//! @param epoch Optional one-element int64 step counter, incremented once. +void minimaxM3Sm107Tp4StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) +{ + TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), + "minimax_m3_sm107_tp4_step_prologue: window must be a contiguous uint8 CUDA tensor"); + TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, + "minimax_m3_sm107_tp4_step_prologue: window is shorter than its header and max_desc descriptors"); + c10::cuda::CUDAGuard const deviceGuard{window.device()}; + int64_t* epochPtr = nullptr; + if (epoch.has_value()) + { + TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 + && epoch->device() == window.device(), + "minimax_m3_sm107_tp4_step_prologue: epoch must be a one-element int64 tensor on the window's device"); + epochPtr = epoch->data_ptr(); + } + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeStepPrologue( + window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_step_prologue launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4StepPrologue); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp new file mode 100644 index 000000000000..1172d809a817 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp4/swigluOai.h): x bf16 [rows, 2d] -> out bf16 +//! [rows, d], PDL-launched with an early release when `pdl`. +void minimaxM3Sm107Tp4SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp4_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp4_swiglu_oai: x rows must be 4-byte aligned with an even half width"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp4_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSwigluOai( + reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, + at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_swiglu_oai launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SwigluOai); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp new file mode 100644 index 000000000000..bee82725b41f --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp @@ -0,0 +1,236 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h" +#include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h" + +#include +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace ds = tensorrt_llm::kernels::minimaxM3Sm107Tp8::draftStep; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_sm107tp8_draft_step expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107tp8_draft_step requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +void checkInt(torch::Tensor const& t, int64_t n, char const* name) +{ + TORCH_CHECK(t.scalar_type() == at::kInt && t.is_contiguous() && t.numel() >= n, "draft_step: ", name, + " must be contiguous int32 with at least ", n, " elements"); +} + +} // namespace + +//! One fused greedy draft pass over `rows` <= 8 requests: the BF16 head (this rank's fragment-ordered shard wf +//! [N / 16, K / 16, 256], or with excOff / excPos / excW / expBase its 12-bit coding, wf uint8 [N * K * 3 / 2]; +//! global ids vocabOffset..) over x[xRows[b]] (or x[b]), the global argmax across the TP ranks through the multicast +//! slots wsMcPtr / wsUc, then draftOut[b, step] = idsNext[b] = token, posNext[b] = posSrc[xRows[b] or b] + 1 and the +//! KV-length bookkeeping (see draftStep.h). +void minimaxM3Sm107Tp8DraftStep(torch::Tensor const& x, OptTensor const& xRows, int64_t rows, torch::Tensor const& wf, + int64_t vocabOffset, torch::Tensor& keys, torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, + int64_t rank, int64_t world, torch::Tensor& draftOut, int64_t step, torch::Tensor& idsNext, + torch::Tensor const& posSrc, torch::Tensor& posNext, torch::Tensor& kvLens, OptTensor const& numAcc, + OptTensor const& seqLens, int64_t draftLen, int64_t numCtas, bool pdl, OptTensor const& outKeys, + OptTensor const& excOff, OptTensor const& excPos, OptTensor const& excW, int64_t expBase, int64_t tileLo, + int64_t tileHi, int64_t totalCtas, OptTensor const& bnd, OptTensor const& bndCnt, + std::optional> fold, int64_t foldRows, int64_t foldTokensPerStep, bool fixedOrder) +{ + checkArch(x); + for (torch::Tensor const* t : std::initializer_list{ + &wf, &keys, &counters, &wsUc, &draftOut, &idsNext, &posSrc, &posNext, &kvLens}) + { + TORCH_CHECK(t->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + for (OptTensor const* t : std::initializer_list{ + &xRows, &numAcc, &seqLens, &outKeys, &excOff, &excPos, &excW, &bnd, &bndCnt}) + { + TORCH_CHECK(!t->has_value() || (*t)->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "draft_step: x must be contiguous BF16 [rows, K]"); + int const K = static_cast(x.size(1)); + TORCH_CHECK(K % 16 == 0, "draft_step: K must be a multiple of 16"); + TORCH_CHECK(rows >= 1 && rows <= ds::kMaxTokens, "draft_step: 1..", ds::kMaxTokens, " rows"); + bool const coded = excOff.has_value(); + TORCH_CHECK(coded == excPos.has_value() && coded == excW.has_value(), + "draft_step: exc_off, exc_pos and exc_w come together"); + int nTiles = 0; + if (coded) + { + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "draft_step: a coded head needs K a multiple of 64"); + TORCH_CHECK(excOff->scalar_type() == at::kInt && excOff->is_contiguous() && excOff->numel() >= 2, + "draft_step: exc_off must be int32 [tiles + 1]"); + nTiles = static_cast(excOff->numel() - 1); + TORCH_CHECK(wf.scalar_type() == at::kByte && wf.is_contiguous() + && wf.numel() == static_cast(nTiles) * 16 * K * 3 / 2, + "draft_step: a coded wf must be the contiguous uint8 stream of [tiles * 16, K]"); + TORCH_CHECK(excPos->scalar_type() == at::kInt && excPos->is_contiguous(), "draft_step: exc_pos must be int32"); + TORCH_CHECK(excW->scalar_type() == at::kFloat && excW->is_contiguous() && excW->numel() == 2 * excPos->numel(), + "draft_step: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "draft_step: exponent base out of range"); + } + else + { + TORCH_CHECK(wf.scalar_type() == at::kBFloat16 && wf.dim() == 3 && wf.is_contiguous() && wf.size(1) == K / 16 + && wf.size(2) == 256, + "draft_step: wf must be the fragment-ordered BF16 weight [N / 16, K / 16, 256]"); + nTiles = static_cast(wf.size(0)); + } + int64_t const rangeHi = tileHi > 0 ? tileHi : nTiles; + TORCH_CHECK(tileLo >= 0 && tileLo < rangeHi && rangeHi <= nTiles, "draft_step: tile range"); + TORCH_CHECK(nTiles >= 1 && numCtas >= 1 && (rangeHi - tileLo + numCtas - 1) / numCtas <= ds::kMaxTilesPerCta, + "draft_step: too many weight tiles per CTA"); + TORCH_CHECK(totalCtas == 0 || totalCtas >= numCtas, "draft_step: total_ctas"); + TORCH_CHECK(bnd.has_value() == bndCnt.has_value(), "draft_step: bnd and bnd_cnt come together"); + if (bnd.has_value()) + { + TORCH_CHECK(bnd->scalar_type() == at::kFloat && bnd->is_contiguous() && bnd->numel() >= int64_t{nTiles} * 128, + "draft_step: bnd must be zero-filled fp32 [tiles, 32, 4]"); + TORCH_CHECK(bndCnt->scalar_type() == at::kInt && bndCnt->is_contiguous() && bndCnt->numel() >= nTiles, + "draft_step: bnd_cnt must be zero-filled int32 [tiles]"); + // At least one tile per CTA: a tile is then shared by at most two CTAs (two addends in bnd). + TORCH_CHECK(rangeHi - tileLo >= numCtas, "draft_step: an equal-chunk split needs a tile per CTA"); + } + TORCH_CHECK(vocabOffset >= 0 && vocabOffset + 16LL * nTiles <= 0xFFFFFFFFLL, "draft_step: vocab offset"); + if (xRows.has_value()) + { + TORCH_CHECK(xRows->scalar_type() == at::kLong && xRows->is_contiguous() && xRows->numel() >= rows, + "draft_step: x_rows must be contiguous int64 [rows]"); + } + else + { + TORCH_CHECK(x.size(0) >= rows, "draft_step: x must hold a row per request"); + } + TORCH_CHECK(world >= 1 && rank >= 0 && rank < world, "draft_step: rank / world"); + TORCH_CHECK(keys.scalar_type() == at::kLong && keys.numel() >= rows, "draft_step: keys must be int64 [>= rows]"); + TORCH_CHECK(counters.scalar_type() == at::kInt && counters.numel() >= 2, "draft_step: counters must be int32 [2]"); + TORCH_CHECK(wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, + "draft_step: ws_uc must hold [2][8][world] slots"); + TORCH_CHECK(draftOut.scalar_type() == at::kInt && draftOut.dim() == 2 && draftOut.stride(1) == 1 + && draftOut.size(0) >= rows && step >= 0 && step < draftOut.size(1), + "draft_step: draft_out must be int32 [rows, > step] with unit column stride"); + checkInt(idsNext, rows, "ids_next"); + checkInt(posNext, rows, "pos_next"); + checkInt(kvLens, rows, "kv_lens"); + TORCH_CHECK(posSrc.scalar_type() == at::kInt && posSrc.is_contiguous(), "draft_step: pos_src must be int32"); + TORCH_CHECK(numAcc.has_value() == seqLens.has_value(), "draft_step: num_acc and seq_lens come together"); + TORCH_CHECK(!numAcc.has_value() || step == 0, "draft_step: the KV rollback (num_acc) is pass 0's"); + if (numAcc.has_value()) + { + checkInt(*numAcc, rows, "num_acc"); + checkInt(*seqLens, rows, "seq_lens"); + } + if (outKeys.has_value()) + { + TORCH_CHECK(outKeys->scalar_type() == at::kLong && outKeys->is_contiguous() && outKeys->numel() >= rows, + "draft_step: out_keys must be contiguous int64 [rows]"); + } + ds::Params p{}; + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_rows = xRows.has_value() ? reinterpret_cast(xRows->data_ptr()) : nullptr; + p.M = static_cast(rows); + p.wf = reinterpret_cast(wf.data_ptr()); + p.n_tiles = nTiles; + p.vocab_offset = static_cast(vocabOffset); + p.keys = reinterpret_cast(keys.data_ptr()); + p.counters = reinterpret_cast(counters.data_ptr()); + p.ws_uc = reinterpret_cast(wsUc.data_ptr()); + p.ws_mc = reinterpret_cast(wsMcPtr); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.out_keys = outKeys.has_value() ? reinterpret_cast(outKeys->data_ptr()) : nullptr; + p.st = ds::StepArgs{draftOut.data_ptr(), static_cast(draftOut.stride(0)), static_cast(step), + idsNext.data_ptr(), posSrc.data_ptr(), posNext.data_ptr(), kvLens.data_ptr(), + numAcc.has_value() ? numAcc->data_ptr() : nullptr, + seqLens.has_value() ? seqLens->data_ptr() : nullptr, static_cast(draftLen)}; + p.cd = ds::CodedArgs{coded ? excOff->data_ptr() : nullptr, + coded ? reinterpret_cast(excPos->data_ptr()) : nullptr, + coded ? reinterpret_cast(excW->data_ptr()) : nullptr, + static_cast((expBase << 7) | (expBase << 23))}; + p.tile_lo = static_cast(tileLo); + p.tile_hi = static_cast(tileHi); + p.total_ctas = static_cast(totalCtas); + p.bnd = bnd.has_value() ? reinterpret_cast(bnd->data_ptr()) : nullptr; + p.bnd_cnt = bnd.has_value() ? reinterpret_cast(bndCnt->data_ptr()) : nullptr; + if (fold.has_value()) + { + // The 14 fold-tail tensors in buildParams' order (ctrl, slots, src_new, dst_new, src_next, dst_next, + // src_draft, dst_draft, src_lens, dst_lens, ids, draft_ids, pos_off, kv_off): the in-kernel tail's operands, + // packed exactly as the standalone fold tail packs them. + std::vector f = *fold; + TORCH_CHECK(f.size() == 14, "minimax_m3_sm107tp8_draft_step: fold takes 14 tensors, got ", f.size()); + p.fold = minimaxM3Sm107Tp8FoldTailParams::buildParams(f[0], f[1], foldRows, f[2], f[3], f[4], f[5], f[6], f[7], + f[8], f[9], f[10], f[11], f[12], f[13], foldTokensPerStep); + p.hasFold = true; + } + p.fixedOrder = fixedOrder; + check(ds::invokeDraftStep(p, K, static_cast(numCtas), pdl, at::cuda::getCurrentCUDAStream(x.get_device()))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// Mutations through the raw multicast address ws_mc_ptr are not visible to the schema: that buffer is owned by the +// Python workspace and never aliased by the op's tensor arguments. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp8_draft_step(Tensor x, Tensor? x_rows, int rows, Tensor wf, int vocab_offset, Tensor(a!) " + "keys, " + "Tensor(b!) counters, Tensor(c!) ws_uc, int ws_mc_ptr, int rank, int world, Tensor(d!) draft_out, int step, " + "Tensor(e!) ids_next, Tensor pos_src, Tensor(f!) pos_next, Tensor(g!) kv_lens, Tensor? num_acc, " + "Tensor(h!)? seq_lens, int draft_len, int num_ctas, bool pdl, Tensor(i!)? out_keys, Tensor? exc_off=None, " + "Tensor? exc_pos=None, Tensor? exc_w=None, int exp_base=0, int tile_lo=0, int tile_hi=0, int total_ctas=0, " + "Tensor(j!)? bnd=None, Tensor(k!)? bnd_cnt=None" + ", Tensor(l!)[]? fold=None, int fold_rows=0, int fold_tokens_per_step=0, bool fixed_order=False) " + "-> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp8_draft_step", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8DraftStep); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp new file mode 100644 index 000000000000..ec93068fc374 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp8::drafterRopeFp8Kv; + +//! The MiniMax-M3 TP8 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h): RoPE on +//! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages +//! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. +void minimaxM3Sm107Tp8DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, + torch::Tensor const& seqLens, std::optional const& specPositionOffsets, torch::Tensor& kvPages, + torch::Tensor const& blockTables, torch::Tensor& qOut, int64_t qLen, int64_t numHeadsQ, int64_t numHeadsKv, + int64_t rotaryDim, int64_t tokensPerBlock, int64_t ropeMode, bool pdl, bool maxSharedCarveout) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + 2 * numHeadsKv) * drk::kHeadDim, "qkv is not [tokens, (Hq + 2 Hkv) * 128]"); + TORCH_CHECK(qLen >= 1 && qkv.size(0) % qLen == 0, "the rows are not whole passes of q_len rows"); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= drk::kHeadDim && rotaryDim % 16 == 0, "unsupported rotary_dim"); + TORCH_CHECK(cosSin.is_cuda() && cosSin.is_contiguous() && cosSin.scalar_type() == at::kFloat + && cosSin.numel() % rotaryDim == 0, + "cos_sin must be the fp32 [positions, rotary_dim / 2] (cos, sin) table"); + int64_t const requests = qkv.size(0) / qLen; + TORCH_CHECK(seqLens.is_cuda() && seqLens.scalar_type() == at::kInt && seqLens.numel() >= requests); + if (specPositionOffsets.has_value()) + { + TORCH_CHECK(specPositionOffsets->is_cuda() && specPositionOffsets->is_contiguous() + && specPositionOffsets->scalar_type() == at::kInt && specPositionOffsets->numel() >= requests * qLen); + } + TORCH_CHECK(kvPages.is_cuda() && kvPages.is_contiguous() && kvPages.dim() == 4 + && kvPages.scalar_type() == at::kFloat8_e4m3fn && kvPages.size(1) == numHeadsKv + && kvPages.size(2) == tokensPerBlock && kvPages.size(3) == drk::kHeadDim, + "kv_pages must be the E4M3 [sub-pages, Hkv, tokens_per_block, 128] cache"); + TORCH_CHECK(blockTables.is_cuda() && blockTables.scalar_type() == at::kInt && blockTables.dim() == 3 + && blockTables.size(0) >= requests && blockTables.size(1) == 2 && blockTables.stride(2) == 1, + "block_tables must be the int32 [requests, 2, max blocks] K / V sub-page table"); + TORCH_CHECK(qOut.is_cuda() && qOut.is_contiguous() && qOut.scalar_type() == at::kFloat8_e4m3fn + && qOut.numel() == qkv.size(0) * numHeadsQ * drk::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{qkv.device()}; + drk::Params params{}; + params.qkv = reinterpret_cast<__nv_bfloat16 const*>(qkv.data_ptr()); + params.cosSin = reinterpret_cast(cosSin.data_ptr()); + params.seqLens = seqLens.data_ptr(); + params.specPositionOffsets = specPositionOffsets.has_value() ? specPositionOffsets->data_ptr() : nullptr; + params.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + params.pool = reinterpret_cast<__nv_fp8_e4m3*>(kvPages.data_ptr()); + params.blockTable = blockTables.data_ptr(); + params.blockTableRequestStride = blockTables.stride(0); + params.blockTableKvStride = blockTables.stride(1); + params.numTokens = static_cast(qkv.size(0)); + params.qLen = static_cast(qLen); + params.numHeadsQ = static_cast(numHeadsQ); + params.numHeadsKv = static_cast(numHeadsKv); + params.rotaryDim = static_cast(rotaryDim); + params.tokensPerBlock = static_cast(tokensPerBlock); + params.ropeMode = static_cast(ropeMode); + cudaError_t const err + = drk::launch(params, pdl, maxSharedCarveout, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp8_drafter_rope_fp8_kv: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp8_drafter_rope_fp8_kv(Tensor qkv, Tensor cos_sin, Tensor seq_lens, " + "Tensor? spec_position_offsets, Tensor(a!) kv_pages, Tensor block_tables, Tensor(b!) q_out, int q_len, " + "int num_heads_q, int num_heads_kv, int rotary_dim, int tokens_per_block, int rope_mode, bool pdl, bool " + "max_shared_carveout=True) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp8_drafter_rope_fp8_kv", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8DrafterRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp new file mode 100644 index 000000000000..e4d5ca2f6b54 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp @@ -0,0 +1,89 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h" +#include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The speculative step's forward-tail store and next-step staging in one launch (the C++ form of the Triton +//! kernels/m3_fold_tail.py, the same writes). The caller checks the row counts against the buffers (fold_tail). +static void foldTailImpl(bool oneWarp, torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + c10::cuda::CUDAGuard const deviceGuard{ctrl.device()}; + auto const p = minimaxM3Sm107Tp8FoldTailParams::buildParams(ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, + srcDraft, dstDraft, srcLens, dstLens, ids, draftIds, posOff, kvOff, tokensPerStep); + auto const stream = at::cuda::getCurrentCUDAStream(ctrl.get_device()); + auto const err = oneWarp + ? tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeFoldTailOneWarp(p, stream) + : tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeFoldTail(p, static_cast(rows), stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_fold_tail launch failed: ", cudaGetErrorString(err)); +} + +//! The tail as its standalone kernel (one CTA per row). +void minimaxM3Sm107Tp8FoldTail(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + foldTailImpl(false, ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, srcDraft, dstDraft, srcLens, dstLens, ids, + draftIds, posOff, kvOff, tokensPerStep); +} + +//! The tail's in-kernel form (one warp over the rows, foldTailWarp), the oracle of a caller that folds it in. +void minimaxM3Sm107Tp8FoldTailOneWarp(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + foldTailImpl(true, ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, srcDraft, dstDraft, srcLens, dstLens, ids, + draftIds, posOff, kvOff, tokensPerStep); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp8_fold_tail(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) dst_new, " + "Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); + m.def( + "minimax_m3_sm107_tp8_fold_tail_one_warp(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) " + "dst_new, Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_fold_tail", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8FoldTail); + m.impl("minimax_m3_sm107_tp8_fold_tail_one_warp", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8FoldTailOneWarp); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h new file mode 100644 index 000000000000..cc5cde28c9d9 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace minimaxM3Sm107Tp8FoldTailParams +{ + +constexpr int64_t kMaxWidth = 32; + +inline void check2d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 2 && t.size(1) <= kMaxWidth, + "minimax_m3_sm107_tp8_fold_tail: ", name, " must be a 2-D int32 CUDA view at most 32 wide"); +} + +inline void check1d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp8_fold_tail: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +//! The fold tail's operands checked and packed (the standalone fold tail op and any op that runs the tail in-kernel, +//! e.g. the drafter's last-pass draft_step, build them identically). +inline tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams buildParams(torch::Tensor const& ctrl, + torch::Tensor const& slots, int64_t rows, torch::Tensor const& srcNew, torch::Tensor& dstNew, + torch::Tensor const& srcNext, torch::Tensor& dstNext, torch::Tensor const& srcDraft, torch::Tensor& dstDraft, + torch::Tensor const& srcLens, torch::Tensor& dstLens, torch::Tensor& ids, torch::Tensor& draftIds, + torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + TORCH_CHECK(ctrl.is_cuda() && ctrl.scalar_type() == torch::kInt32 && ctrl.numel() >= 1, + "minimax_m3_sm107_tp8_fold_tail: ctrl must be an int32 CUDA tensor"); + TORCH_CHECK(slots.is_cuda() && slots.scalar_type() == torch::kInt64 && slots.dim() == 1 && slots.stride(0) == 1, + "minimax_m3_sm107_tp8_fold_tail: slots must be a contiguous 1-D int64 CUDA tensor"); + TORCH_CHECK(rows > 0 && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, + "minimax_m3_sm107_tp8_fold_tail: rows must be positive and tokens_per_step in [1, 32]"); + check2d(srcNew, "src_new"); + check2d(dstNew, "dst_new"); + check2d(srcNext, "src_next"); + check2d(dstNext, "dst_next"); + check2d(srcDraft, "src_draft"); + check2d(dstDraft, "dst_draft"); + check1d(srcLens, "src_lens"); + check1d(dstLens, "dst_lens"); + check1d(ids, "ids"); + check1d(draftIds, "draft_ids"); + check1d(posOff, "pos_off"); + check1d(kvOff, "kv_off"); + tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams p{}; + p.ctrl = ctrl.data_ptr(); + p.slots = slots.data_ptr(); + p.srcNew = srcNew.data_ptr(); + p.sn0 = srcNew.stride(0); + p.sn1 = srcNew.stride(1); + p.snW = static_cast(srcNew.size(1)); + p.dstNew = dstNew.data_ptr(); + p.dn0 = dstNew.stride(0); + p.dn1 = dstNew.stride(1); + p.dnW = static_cast(dstNew.size(1)); + p.srcNext = srcNext.data_ptr(); + p.sx0 = srcNext.stride(0); + p.sx1 = srcNext.stride(1); + p.sxW = static_cast(srcNext.size(1)); + p.dstNext = dstNext.data_ptr(); + p.dx0 = dstNext.stride(0); + p.dx1 = dstNext.stride(1); + p.dxW = static_cast(dstNext.size(1)); + p.srcDraft = srcDraft.data_ptr(); + p.sd0 = srcDraft.stride(0); + p.sd1 = srcDraft.stride(1); + p.sdW = static_cast(srcDraft.size(1)); + p.dstDraft = dstDraft.data_ptr(); + p.dd0 = dstDraft.stride(0); + p.dd1 = dstDraft.stride(1); + p.ddW = static_cast(dstDraft.size(1)); + p.srcLens = srcLens.data_ptr(); + p.dstLens = dstLens.data_ptr(); + p.ids = ids.data_ptr(); + p.draftIds = draftIds.data_ptr(); + p.posOff = posOff.data_ptr(); + p.kvOff = kvOff.data_ptr(); + p.tokensPerStep = static_cast(tokensPerStep); + return p; +} + +} // namespace minimaxM3Sm107Tp8FoldTailParams +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp new file mode 100644 index 000000000000..29b9a0ab391b --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107Tp8::fullCtxReduce; + +//! Merges the PrimTS split-KV decode scratch (partial_o [B, Hkv, S, R, 128] 16-bit, partial_stats [B, Hkv, S, R] +//! fp32 log2-LSE) of the MiniMax-M3 TP8 full-context decode into out (BF16 [B * SQ, Hkv * hR, 128]); with q8 / sf +//! also its MXFP8 copy (E4M3 [>= B * SQ, Hkv * hR * 128] + 128x4-swizzled UE8M0 scales, as mxfp8_quantize(out, True)). +void minimaxM3Sm107Tp8FullCtxReduce(torch::Tensor const& partialO, torch::Tensor const& partialStats, + torch::Tensor const& seqLensKv, torch::Tensor& out, int64_t seqLenQ, int64_t qTokensPerCta, int64_t tileKv, + int64_t instsKv, bool causal, bool pdl, std::optional const& q8, + std::optional const& sf) +{ + TORCH_CHECK(partialO.is_cuda() && partialO.is_contiguous() && partialO.dim() == 5); + TORCH_CHECK(partialO.scalar_type() == at::kHalf || partialO.scalar_type() == at::kBFloat16); + TORCH_CHECK(partialO.size(4) == fcr::kHeadDim && partialO.size(2) <= fcr::kMaxSplits); + TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() + && partialStats.sizes() == partialO.sizes().slice(0, 4)); + TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + int const rows = static_cast(partialO.size(3)); + TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); + int const headsPerKv = rows / static_cast(seqLenQ); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.is_contiguous() + && out.numel() == partialO.size(0) * seqLenQ * partialO.size(1) * headsPerKv * fcr::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{partialO.device()}; + fcr::Params params{partialO.data_ptr(), partialStats.data_ptr(), seqLensKv.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), static_cast(partialO.size(0)), + static_cast(partialO.size(1)), static_cast(partialO.size(2)), rows, static_cast(seqLenQ), + headsPerKv, static_cast(qTokensPerCta), static_cast(tileKv), static_cast(instsKv), causal, + partialO.scalar_type() == at::kBFloat16}; + TORCH_CHECK(q8.has_value() == sf.has_value()); + if (q8.has_value()) + { + int64_t const width = partialO.size(1) * headsPerKv * fcr::kHeadDim; + int64_t const tokens = partialO.size(0) * seqLenQ; + int64_t const tileRows = (tokens + 127) / 128 * 128; + TORCH_CHECK(q8->is_cuda() && q8->is_contiguous() && q8->scalar_type() == at::kFloat8_e4m3fn && q8->dim() == 2 + && q8->size(1) == width && q8->size(0) >= tokens); + TORCH_CHECK(sf->is_cuda() && sf->is_contiguous() && sf->scalar_type() == at::kByte + && sf->numel() >= tileRows * (width / 32)); + params.q8 = reinterpret_cast(q8->data_ptr()); + params.sf = sf->data_ptr(); + params.numSfKTiles = static_cast(width / 128); + } + cudaError_t const err + = fcr::invokeFullCtxReduce(params, pdl, at::cuda::getCurrentCUDAStream(partialO.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp8_fullctx_reduce: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp8_fullctx_reduce(Tensor partial_o, Tensor partial_stats, Tensor seq_lens_kv, " + "Tensor(a!) out, int seq_len_q, int q_tokens_per_cta, int tile_kv, int insts_kv, bool causal, bool pdl, " + "Tensor(b!)? q8=None, Tensor(c!)? sf=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp8_fullctx_reduce", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8FullCtxReduce); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp new file mode 100644 index 000000000000..eb3759b28a58 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp @@ -0,0 +1,181 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h" + +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxTokens = 48; +constexpr int64_t kCluster = 8; +constexpr int64_t kMaxWorld = 8; + +} // namespace + +//! MiniMax-M3 SM107 TP8 multicast all-reduce (stock rounding points: BF16 rank sum, then the residual add) + residual + +//! RMSNorm (+ optional MoE finalize, + optional MXFP8 output). +//! +//! @param x BF16 [M, H] partial, or with `perm` the permuted FC2 output [rows, H]. +//! @param x2 Optional BF16 [M, H] second addend (the shared expert in finalize mode). +//! @param residual BF16 [M, H], 1 <= M <= 48. +//! @param normWeight BF16 [H]. +//! @param wsUc This rank's unicast view of the exchange workspace (int32 words); re-armed in place. +//! @param wsMcPtr Device address of the workspace's multicast mapping. +//! @param flags Int32 per-CTA call-parity counters, at least maxt * 8; updated in place. +//! @param cluster Cluster size; 8 is the only one built. +//! @param perm, scales Finalize mode: int32 [M, topk] permuted-row index and fp32 / bf16 [M, topk] routing weights. +//! @param maxt Token capacity of the workspace bucket (8, 16 or 48). +//! @param quant Also return the MXFP8 copy of the normed rows (E4M3 [M, H], swizzled UE8M0 scales). +//! @param prepushed 0, or 2 / 3 when producers already pushed this rank's rows (see the kernel header). +//! @param opts Epilogue variant; every variant is bit-identical. +//! @param capOut Optional BF16 [>= M, H] view (row stride a multiple of 8, 16-byte aligned) that receives the Eagle3 +//! capture of the updated residual: m3_spec_tail capture_norm with capW [H] and capEps, bit for bit (H 6144; not the +//! drafter boundary, opts bit 16). +//! @return [normed, updated residual] or, with `quant`, [normed, updated residual, E4M3 rows, scale bytes]. +std::vector minimaxM3Sm107Tp8McAllreduceNorm(torch::Tensor const& x, + std::optional const& x2, torch::Tensor const& residual, torch::Tensor const& normWeight, double eps, + torch::Tensor& wsUc, int64_t wsMcPtr, torch::Tensor& flags, int64_t rank, int64_t world, int64_t cluster, bool pdl, + std::optional const& perm, std::optional const& scales, int64_t maxt, bool quant, + int64_t prepushed, int64_t opts, std::optional const& capOut, + std::optional const& capW, double capEps) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 107, "minimax_m3_sm107_tp8_mc_allreduce_norm requires SM107, got SM", sm); + TORCH_CHECK(cluster == kCluster, "minimax_m3_sm107_tp8_mc_allreduce_norm: only cluster 8 is built, got ", cluster); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107Tp8McAllreduceNormSupported( + static_cast(maxt), static_cast(opts)), + "minimax_m3_sm107_tp8_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); + TORCH_CHECK( + world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); + + tensorrt_llm::kernels::MinimaxM3Sm107Tp8McAllreduceNormParams p; + if (x2.has_value()) + { + TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), + "x2 must be a contiguous bf16 tensor shaped like residual"); + p.x2 = reinterpret_cast<__nv_bfloat16 const*>(x2->data_ptr()); + } + if (perm.has_value()) + { + // Finalize mode: x is the permuted FC2 output [rows, H]. + TORCH_CHECK(scales.has_value() && perm->scalar_type() == at::kInt && perm->is_contiguous() && perm->dim() == 2 + && perm->size(0) == residual.size(0), + "perm must be a contiguous int32 [M, topk] tensor, with scales"); + TORCH_CHECK(scales->is_contiguous() && scales->numel() == perm->numel() + && (scales->scalar_type() == at::kFloat || scales->scalar_type() == at::kBFloat16), + "scales must be a contiguous fp32 or bf16 tensor with one weight per perm entry"); + TORCH_CHECK(x.dim() == 2 && x.size(1) == residual.size(1), "fc2 must be [rows, H]"); + p.fc2 = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.perm = perm->data_ptr(); + p.scales = scales->data_ptr(); + p.scaleFp32 = scales->scalar_type() == at::kFloat; + p.topk = static_cast(perm->size(1)); + } + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2, + "x must be a contiguous 2-D bf16 CUDA tensor"); + TORCH_CHECK(residual.dim() == 2 && residual.size(0) >= 1 && residual.size(0) <= std::min(kMaxTokens, maxt), + "residual must hold 1..min(48, bucket) rows"); + TORCH_CHECK(residual.scalar_type() == at::kBFloat16 && residual.is_contiguous() + && (perm.has_value() || residual.sizes() == x.sizes()), + "residual must be a contiguous bf16 tensor shaped like x"); + TORCH_CHECK( + normWeight.scalar_type() == at::kBFloat16 && normWeight.numel() == residual.size(1), "norm_w must be bf16 [H]"); + TORCH_CHECK(wsUc.scalar_type() == at::kInt && wsUc.is_cuda(), "ws_uc must be an int32 CUDA tensor"); + TORCH_CHECK(flags.scalar_type() == at::kInt && flags.is_contiguous() && flags.numel() >= maxt * kCluster, + "flags must be a contiguous int32 tensor of at least bucket * 8 counters"); + + int64_t const m = residual.size(0), h = residual.size(1); + TORCH_CHECK(!quant || h % 128 == 0, "MXFP8 output needs H % 128 == 0"); + c10::cuda::CUDAGuard const deviceGuard{residual.device()}; + auto outNormed = torch::empty_like(residual); + auto outResidual = torch::empty_like(residual); + auto outQ = quant ? torch::empty({m, h}, residual.options().dtype(at::kFloat8_e4m3fn)) : torch::Tensor(); + auto outSf = quant ? torch::empty({(m + 127) / 128 * 128 * (h / 32)}, residual.options().dtype(at::kByte)) + : torch::Tensor(); + + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.numTokens = static_cast(m); + p.hidden = static_cast(h); + p.residual = reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()); + p.normWeight = reinterpret_cast<__nv_bfloat16 const*>(normWeight.data_ptr()); + p.eps = static_cast(eps); + p.wsUc = wsUc.data_ptr(); + p.wsMc = reinterpret_cast(wsMcPtr); + p.flags = flags.data_ptr(); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.outNormed = reinterpret_cast<__nv_bfloat16*>(outNormed.data_ptr()); + p.outResidual = reinterpret_cast<__nv_bfloat16*>(outResidual.data_ptr()); + p.outQuant = quant ? reinterpret_cast(outQ.data_ptr()) : nullptr; + p.outSf = quant ? outSf.data_ptr() : nullptr; + p.prepushed = static_cast(prepushed); + if (capOut.has_value()) + { + TORCH_CHECK((opts & 16) == 0, "the capture epilogue is not built for the drafter boundary (opts bit 16)"); + TORCH_CHECK(capOut->scalar_type() == at::kBFloat16 && capOut->dim() == 2 && capOut->size(0) >= m + && capOut->size(1) == h && capOut->stride(1) == 1 && capOut->stride(0) % 8 == 0 + && reinterpret_cast(capOut->data_ptr()) % 16 == 0 && capOut->device() == residual.device(), + "cap_out must be a bf16 [>= M, H] view with unit column stride, a row stride that is a multiple of 8 and " + "16-byte alignment"); + TORCH_CHECK( + capW.has_value() && capW->scalar_type() == at::kBFloat16 && capW->is_contiguous() && capW->numel() == h, + "cap_w must be a contiguous bf16 [H] tensor"); + p.capOut = reinterpret_cast<__nv_bfloat16*>(capOut->data_ptr()); + p.capStride = capOut->stride(0); + p.capWeight = reinterpret_cast<__nv_bfloat16 const*>(capW->data_ptr()); + p.capEps = static_cast(capEps); + } + tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp8McAllreduceNorm(p, static_cast(maxt), + static_cast(opts), pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); + if (quant) + { + return {outNormed, outResidual, outQ, outSf}; + } + return {outNormed, outResidual}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp8_mc_allreduce_norm(Tensor x, Tensor? x2, Tensor residual, Tensor norm_w, float eps, " + "Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) flags, int rank, int world, int cluster, bool pdl, " + "Tensor? perm, Tensor? scales, int maxt, bool quant, int prepushed, int opts, Tensor(c!)? cap_out=None, " + "Tensor? cap_w=None, float cap_eps=0.) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_mc_allreduce_norm", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8McAllreduceNorm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp new file mode 100644 index 000000000000..acd7061eadbb --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp @@ -0,0 +1,185 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp8::producerPw8; + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPageSize = 128; + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "sm107_tp8_producer_pw: ", name, " must be a CUDA tensor"); +} + +} // namespace + +//! The SM107 TP8 NVFP4 QKV + indexer norm/RoPE/KV-insert producer (same arguments and outputs as +//! minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert at 8 query, 1 KV and 1 index head, head_dim 128, rotary_dim 64): +//! returns (q FP8 [T, 8, 128], index_q FP8 [T, 1, 128]) and writes the NVFP4 K/V of the KV head and the FP8 index-K +//! of each token's cache slot. `early` triggers the dependents on entry; `prewait_step_inputs` reads the position ids, +//! the slots and the RoPE rows before the grid-dependency wait (none of them may be written by a PDL predecessor). +std::vector minimaxM3Sm107Tp8ProducerPw(torch::Tensor const& packed, torch::Tensor& kvData, + torch::Tensor& kvScale, torch::Tensor& indexK, torch::Tensor const& outCacheLoc, torch::Tensor const& kvQuantScale, + double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds, bool pdl, + bool early, bool prewaitStepInputs) +{ + checkCuda(packed, "packed"); + for (torch::Tensor const* t : std::initializer_list{&kvData, &kvScale, &indexK, &outCacheLoc, + &kvQuantScale, &qWeight, &kWeight, &indexQWeight, &indexKWeight, &rotaryCosSin, &positionIds}) + { + TORCH_CHECK(t->is_cuda() && t->get_device() == packed.get_device(), + "sm107_tp8_producer_pw: every operand must be on packed's CUDA device"); + } + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0 + && reinterpret_cast(kvData.data_ptr()) % 2 == 0 + && reinterpret_cast(indexK.data_ptr()) % 4 == 0, + "sm107_tp8_producer_pw: packed and the caches must start at the vector-load / store alignment"); + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp8_producer_pw requires SM100 or newer"); + int64_t const headsQ = pw::kHeadsQ; + int64_t const headsKV = pw::kHeadsKV; + int64_t const headsIndex = pw::kHeadsIndex; + int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; + TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim + && packed.is_contiguous(), + "sm107_tp8_producer_pw: packed must be contiguous BF16 [T, 12 * 128]"); + // The kernel's layout (the native producer's checkMinimaxM3Nvfp4HndKVPools at one KV head): HND pools + // [pages, 2 planes, 1 head, 128 tokens, 64 data / 8 scale bytes], dense token rows and head blocks, planes and + // pages that do not overlap. + TORCH_CHECK(kvData.scalar_type() == at::kByte && kvData.dim() == 5 && kvData.size(0) > 0 && kvData.size(1) == 2 + && kvData.size(2) == headsKV && kvData.size(3) == kPageSize && kvData.size(4) == kHeadDim / 2, + "sm107_tp8_producer_pw: kv_data must be the NVFP4 pool [pages, 2, 1, 128, 64] uint8"); + TORCH_CHECK(kvScale.scalar_type() == at::kByte && kvScale.dim() == 5 && kvScale.size(0) == kvData.size(0) + && kvScale.size(1) == 2 && kvScale.size(2) == headsKV && kvScale.size(3) == kPageSize + && kvScale.size(4) == kHeadDim / 16, + "sm107_tp8_producer_pw: kv_scale must be the E4M3 scale pool [pages, 2, 1, 128, 8] uint8"); + TORCH_CHECK(kvData.stride(4) == 1 && kvData.stride(3) == kvData.size(4) + && kvData.stride(2) == kvData.size(3) * kvData.stride(3), + "sm107_tp8_producer_pw: kv_data must have dense token rows and head blocks"); + TORCH_CHECK(kvScale.stride(4) == 1 && kvScale.stride(3) == kvScale.size(4) + && kvScale.stride(2) == kvScale.size(3) * kvScale.stride(3), + "sm107_tp8_producer_pw: kv_scale must have dense token rows and head blocks"); + TORCH_CHECK(kvData.stride(1) >= kvData.size(2) * kvData.stride(2) && kvData.stride(0) >= 2 * kvData.stride(1) + && kvScale.stride(1) >= kvScale.size(2) * kvScale.stride(2) && kvScale.stride(0) >= 2 * kvScale.stride(1), + "sm107_tp8_producer_pw: pool planes or pages overlap"); + TORCH_CHECK(kvData.stride(0) % 2 == 0 && kvData.stride(1) % 2 == 0, + "sm107_tp8_producer_pw: kv_data page and plane strides must keep the uint16 stores aligned"); + TORCH_CHECK(indexK.scalar_type() == at::kFloat8_e4m3fn && indexK.dim() == 4 && indexK.size(0) == kvData.size(0) + && indexK.size(1) == 1 && indexK.size(2) == kPageSize && indexK.size(3) == kHeadDim && indexK.stride(3) == 1 + && indexK.stride(2) == kHeadDim && indexK.stride(0) >= kPageSize * kHeadDim && indexK.stride(0) % 4 == 0, + "sm107_tp8_producer_pw: index_k must be FP8 [pages, 1, 128, 128] with a page stride keeping the 32-bit stores " + "aligned"); + int64_t const numTokens = packed.size(0); + TORCH_CHECK( + outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous() && outCacheLoc.numel() >= numTokens, + "sm107_tp8_producer_pw: out_cache_loc must be int32 [>= T]"); + TORCH_CHECK( + positionIds.scalar_type() == at::kInt && positionIds.is_contiguous() && positionIds.numel() >= numTokens, + "sm107_tp8_producer_pw: position_ids must be int32 [>= T]"); + TORCH_CHECK(kvQuantScale.scalar_type() == at::kFloat && kvQuantScale.is_contiguous() && kvQuantScale.numel() >= 3, + "sm107_tp8_producer_pw: kv_quant_scale must be fp32 [>= 3]"); + TORCH_CHECK(rotaryCosSin.scalar_type() == at::kFloat && rotaryCosSin.is_contiguous() && rotaryCosSin.dim() == 3 + && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == 32 + && reinterpret_cast(rotaryCosSin.data_ptr()) % 16 == 0, + "sm107_tp8_producer_pw: rotary_cos_sin must be fp32 [positions, 2, 32], 16-byte aligned (read as float4)"); + for (auto const* w : {&qWeight, &kWeight, &indexQWeight, &indexKWeight}) + { + TORCH_CHECK(w->scalar_type() == at::kBFloat16 && w->is_contiguous() && w->numel() == kHeadDim + && reinterpret_cast(w->data_ptr()) % 8 == 0, + "sm107_tp8_producer_pw: norm weights must be BF16 [128], 8-byte aligned (read as 4-value words)"); + } + auto const fp8 = packed.options().dtype(at::kFloat8_e4m3fn); + auto qOut = torch::empty({numTokens, headsQ, kHeadDim}, fp8); + auto indexQOut = torch::empty({numTokens, headsIndex, kHeadDim}, fp8); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + pw::Params p{}; + p.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + p.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + p.indexQOut = reinterpret_cast<__nv_fp8_e4m3*>(indexQOut.data_ptr()); + p.kvData = kvData.data_ptr(); + p.kvScale = kvScale.data_ptr(); + p.indexK = reinterpret_cast<__nv_fp8_e4m3*>(indexK.data_ptr()); + p.outCacheLoc = outCacheLoc.data_ptr(); + p.kvQuantScale = kvQuantScale.data_ptr(); + p.dataPageStride = kvData.stride(0); + p.dataPlaneStride = kvData.stride(1); + p.dataHeadStride = kvData.stride(2); + p.dataTokenStride = kvData.stride(3); + p.scalePageStride = kvScale.stride(0); + p.scalePlaneStride = kvScale.stride(1); + p.scaleHeadStride = kvScale.stride(2); + p.indexPageStride = indexK.stride(0); + p.indexTokenStride = indexK.stride(2); + p.numPages = kvData.size(0); + p.numTokens = static_cast(numTokens); + p.eps = static_cast(eps); + p.qWeight = reinterpret_cast<__nv_bfloat16 const*>(qWeight.data_ptr()); + p.kWeight = reinterpret_cast<__nv_bfloat16 const*>(kWeight.data_ptr()); + p.indexQWeight = reinterpret_cast<__nv_bfloat16 const*>(indexQWeight.data_ptr()); + p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); + p.rotaryCosSin = rotaryCosSin.data_ptr(); + p.positionIds = positionIds.data_ptr(); + cudaError_t const err = pw::invokeProducerPw8(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "sm107_tp8_producer_pw launch: ", cudaGetErrorString(err)); + } + return {qOut, indexQOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp8_producer_pw(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) index_k, " + "Tensor out_cache_loc, Tensor kv_quant_scale, float eps, Tensor q_weight, Tensor k_weight, " + "Tensor index_q_weight, Tensor index_k_weight, Tensor rotary_cos_sin, Tensor position_ids, bool pdl, " + "bool early, bool prewait_step_inputs=False) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_producer_pw", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8ProducerPw); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp new file mode 100644 index 000000000000..f6200c9aa906 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp @@ -0,0 +1,113 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3Sm107Tp8::sampledAccept; + +//! uint64 words of the sampled-accept exchange workspace for `world` ranks and SA-1R variant `oneRoundK`. +int64_t minimaxM3Sm107Tp8SampledAcceptWsWords(int64_t world, int64_t oneRoundK) +{ + return sa::sampledAcceptWsWords(static_cast(world), static_cast(oneRoundK)); +} + +//! Sampled (temperature + top-p) verification from this rank's BF16 logits shard [rows, Vs], rows = nreq * +//! (draft_len + 1) <= 48: the sampled token of every row (tokens), then the strict acceptance against draft (accTokens, +//! numAcc, gatherIds, optional nextNew column 0). Per-row Philox (seeds, offsets); the exchange goes through the +//! multicast workspace wsMcPtr / wsUc with the (finished rows, call) counters. +void minimaxM3Sm107Tp8SampledAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor const& temperatures, + torch::Tensor const& topPs, torch::Tensor const& seeds, torch::Tensor const& offsets, torch::Tensor& wsUc, + int64_t wsMcPtr, torch::Tensor& counters, torch::Tensor& tokens, int64_t rank, int64_t world, + torch::Tensor const& draft, torch::Tensor const& batchIdx, torch::Tensor& accTokens, torch::Tensor& numAcc, + torch::Tensor& gatherIds, int64_t forceTotal, std::optional const& nextNew, + std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, + bool prewaitParams, bool earlyTrigger, int64_t oneRoundK, std::optional const& decided, bool ilpWalk) +{ + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + sa::Args a; + a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + a.ld = logits.stride(0); + a.rows = static_cast(logits.size(0)); + a.Vs = static_cast(logits.size(1)); + TORCH_CHECK(a.Vs <= sa::kMaxVs && a.rows <= sa::kMaxRows && world <= sa::kMaxWorld + && static_cast(vocabOffset) + a.Vs <= (1LL << sa::kIdBits)); + a.vocab_offset = static_cast(vocabOffset); + a.temperatures = temperatures.data_ptr(); + a.top_ps = topPs.data_ptr(); + a.seeds = reinterpret_cast(seeds.data_ptr()); + a.offsets = reinterpret_cast(offsets.data_ptr()); + a.ws_uc = reinterpret_cast(wsUc.data_ptr()); + a.ws_mc = reinterpret_cast(wsMcPtr); + a.counters = reinterpret_cast(counters.data_ptr()); + a.tokens = tokens.data_ptr(); + a.extra_attempts = extraAttempts.has_value() ? extraAttempts->data_ptr() : nullptr; + a.rank = static_cast(rank); + a.world = static_cast(world); + a.draft = draft.data_ptr(); + a.batch_idx = batchIdx.data_ptr(); + a.acc_tokens = accTokens.data_ptr(); + a.num_acc = numAcc.data_ptr(); + a.gather_ids = reinterpret_cast(gatherIds.data_ptr()); + a.draft_len = static_cast(draft.size(1)); + a.nreq = static_cast(draft.size(0)); + a.force_total = static_cast(forceTotal); + a.next_new = nextNew.has_value() ? nextNew->data_ptr() : nullptr; + a.next_stride = nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0; + a.force_vec = forceVec.has_value() ? forceVec->data_ptr() : nullptr; + a.prewait_params = prewaitParams ? 1 : 0; + a.early_trigger = earlyTrigger ? 1 : 0; + TORCH_CHECK(oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "one_round_k must be 0, 9 or 82, got ", oneRoundK); + a.one_round_k = static_cast(oneRoundK); + a.decided = decided.has_value() ? decided->data_ptr() : nullptr; + a.ilp_walk = ilpWalk ? 1 : 0; + cudaError_t const err = sa::invokeSampledAccept(a, usePdl, at::cuda::getCurrentCUDAStream(logits.get_device())); + TORCH_CHECK(err == cudaSuccess, "m3_sampled_accept launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// ws_mc_ptr is the multicast address of ws_uc's mapping: the op writes through it, which the schema cannot show. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + // No tensor arguments, so no dispatch key: registered with the schema as a catch-all kernel. + m.def("minimax_m3_sm107tp8_sampled_accept_ws_words(int world, int one_round_k=0) -> int", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SampledAcceptWsWords); + m.def( + "minimax_m3_sm107tp8_sampled_accept(Tensor logits, int vocab_offset, Tensor temperatures, Tensor top_ps, " + "Tensor seeds, Tensor offsets, Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) counters, Tensor(c!) tokens, " + "int rank, int world, Tensor draft, Tensor batch_idx, Tensor(d!) acc_tokens, Tensor(e!) num_acc, " + "Tensor(f!) gather_ids, int force_total, Tensor(g!)? next_new, Tensor? force_vec, " + "Tensor(h!)? extra_attempts, bool use_pdl, bool prewait_params=True, bool early_trigger=False, " + "int one_round_k=0, Tensor(i!)? decided=None, bool ilp_walk=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp8_sampled_accept", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SampledAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp new file mode 100644 index 000000000000..a6cda645ba7f --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp8/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, +//! d]. +void minimaxM3Sm107Tp8SiluAndMul(torch::Tensor const& x, torch::Tensor& out) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp8_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp8_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp8_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); + TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp8_silu_and_mul: at most 65535 rows"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err + = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), + reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_silu_and_mul launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SiluAndMul); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp new file mode 100644 index 000000000000..dc80380b48bb --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp @@ -0,0 +1,189 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeads = 8; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kTopK = 16; + +void checkBytes(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, "minimax_m3_sm107_tp8_sparse_decode_lc8: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3 TP8 long-context sparse decode over the selected NVFP4 blocks (8 query heads on one KV head, +//! head_dim 128, page 128, top-k 16). Dependents are released right after the grid dependency wait. +//! +//! @param q E4M3 bytes [T, 8, 128], unit stride in the last dim. +//! @param k, v Packed E2M1 bytes [pages, 1, 128, 64] (the per-layer HND views). +//! @param kBlockScale, vBlockScale E4M3 bytes [pages, 1, 128, 8]: K token-major, V in 4x4 token-quad order. +//! @param kGlobalScale, vGlobalScale One-element float32 per-tensor dequant scales. +//! @param topk Int32 [T, 16] ascending selected block ids with -1 tail (kv head 0 of the head-major table). +//! @param out BF16 [T, 1024], written. q8 / sf, when given, receive the MXFP8 copy and its swizzled UE8M0 scales. +//! @param prewait Stage seq_lens, q and the request's block-table row before the grid dependency wait (same output). +void minimaxM3Sm107Tp8SparseDecodeLc8(torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& kBlockScale, torch::Tensor const& vBlockScale, torch::Tensor const& kGlobalScale, + torch::Tensor const& vGlobalScale, torch::Tensor const& topk, torch::Tensor const& blockTable, + torch::Tensor const& seqLens, torch::Tensor& out, std::optional q8, std::optional sf, + int64_t numSfKTiles, int64_t decodeQueryLen, double smScale, bool linearSoftmax, bool usePdl, bool prewait, + std::optional prefetchA, std::optional prefetchB, std::optional n1Self, + std::optional n1Next, std::optional n1K, std::optional n1V, + std::optional n1KScale, std::optional n1VScale, bool identity, bool oneDiv, + bool identityStep) +{ + TORCH_CHECK(q.is_cuda(), "minimax_m3_sm107_tp8_sparse_decode_lc8: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m3_sm107_tp8_sparse_decode_lc8 requires an SM100-family GPU"); + checkBytes(q, "q"); + checkBytes(k, "k"); + checkBytes(v, "v"); + checkBytes(kBlockScale, "k_block_scale"); + checkBytes(vBlockScale, "v_block_scale"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeads && q.size(2) == kHeadDim && q.stride(2) == 1, + "minimax_m3_sm107_tp8_sparse_decode_lc8: q must be [T, 8, 128] with unit last-dim stride"); + TORCH_CHECK(kGlobalScale.scalar_type() == torch::kFloat32 && vGlobalScale.scalar_type() == torch::kFloat32, + "minimax_m3_sm107_tp8_sparse_decode_lc8: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m3_sm107_tp8_sparse_decode_lc8: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_sm107_tp8_sparse_decode_lc8: block_table and seq_lens must be int32"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.size(0) == q.size(0) + && out.size(1) == kHeads * kHeadDim && out.stride(1) == 1, + "minimax_m3_sm107_tp8_sparse_decode_lc8: out must be bfloat16 [T, 1024] with unit last-dim stride"); + TORCH_CHECK( + q8.has_value() == sf.has_value(), "minimax_m3_sm107_tp8_sparse_decode_lc8: pass both q8 and sf, or neither"); + bool const mxfp8 = q8.has_value(); + if (mxfp8) + { + checkBytes(*q8, "q8"); + checkBytes(*sf, "sf"); + } + + tensorrt_llm::kernels::minimaxM3Sm107Tp8::MinimaxM3Tp8SparseDecodeLc8Params prm{}; + prm.q = reinterpret_cast(q.data_ptr()); + prm.q_tok = q.stride(0); + prm.q_head = q.stride(1); + prm.k = reinterpret_cast(k.data_ptr()); + prm.v = reinterpret_cast(v.data_ptr()); + prm.k_blk = k.stride(0); + prm.v_blk = v.stride(0); + prm.ksc = reinterpret_cast(kBlockScale.data_ptr()); + prm.vsc = reinterpret_cast(vBlockScale.data_ptr()); + prm.ksc_blk = kBlockScale.stride(0); + prm.vsc_blk = vBlockScale.stride(0); + prm.kgs = kGlobalScale.data_ptr(); + prm.vgs = vGlobalScale.data_ptr(); + prm.topk = topk.data_ptr(); + prm.t_tok = topk.stride(0); + prm.t_k = topk.stride(1); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.seq_lens = seqLens.data_ptr(); + prm.dql = static_cast(decodeQueryLen); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + prm.out_stride = out.stride(0); + prm.q8 = mxfp8 ? reinterpret_cast(q8->data_ptr()) : nullptr; + prm.sf = mxfp8 ? reinterpret_cast(sf->data_ptr()) : nullptr; + prm.num_sf_k_tiles = static_cast(numSfKTiles); + prm.identity = identity ? 1 : 0; + prm.identity_step = (identity || identityStep) ? 1 : 0; + prm.one_div = oneDiv ? 1 : 0; + auto const prefetchRange = [](std::optional const& t, uint8_t const*& ptr, int64_t& bytes) + { + ptr = nullptr; + bytes = 0; + if (t.has_value()) + { + TORCH_CHECK(t->is_cuda() && t->is_contiguous(), + "minimax_m3_sm107_tp8_sparse_decode_lc8: a prefetch range must be a contiguous CUDA tensor"); + ptr = reinterpret_cast(t->data_ptr()); + bytes = t->numel() * static_cast(t->element_size()); + } + }; + prefetchRange(prefetchA, prm.pf0, prm.pf0_bytes); + prefetchRange(prefetchB, prm.pf1, prm.pf1_bytes); + bool const n1 = n1Self.has_value(); + TORCH_CHECK(n1 == n1Next.has_value() && n1 == n1K.has_value() && n1 == n1V.has_value() && n1 == n1KScale.has_value() + && n1 == n1VScale.has_value(), + "minimax_m3_sm107_tp8_sparse_decode_lc8: pass all six next-layer prefetch tensors, or none"); + if (n1) + { + for (auto const* t : {&*n1Self, &*n1Next}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == torch::kInt32 && t->is_contiguous() && t->dim() == 2 + && t->size(1) == kTopK && t->size(0) >= q.size(0), + "minimax_m3_sm107_tp8_sparse_decode_lc8: the id buffers must be contiguous int32 [>= T, 16]"); + } + TORCH_CHECK(n1K->stride(0) == k.stride(0) && n1V->stride(0) == v.stride(0) + && n1KScale->stride(0) == kBlockScale.stride(0) && n1VScale->stride(0) == vBlockScale.stride(0) + && n1K->size(0) == k.size(0), + "minimax_m3_sm107_tp8_sparse_decode_lc8: the next layer's pools must have this layer's page strides"); + prm.n1_self = n1Self->data_ptr(); + prm.n1_next = n1Next->data_ptr(); + prm.n1_k = reinterpret_cast(n1K->data_ptr()); + prm.n1_v = reinterpret_cast(n1V->data_ptr()); + prm.n1_ksc = reinterpret_cast(n1KScale->data_ptr()); + prm.n1_vsc = reinterpret_cast(n1VScale->data_ptr()); + prm.n1_pages = n1K->size(0); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSparseDecodeLc8m( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, prewait, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_sparse_decode_lc8 launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp8_sparse_decode_lc8(Tensor q, Tensor k, Tensor v, Tensor k_block_scale, Tensor " + "v_block_scale, " + "Tensor k_global_scale, Tensor v_global_scale, Tensor topk, Tensor block_table, Tensor seq_lens, " + "Tensor(a!) out, Tensor(b!)? q8, Tensor(c!)? sf, int num_sf_k_tiles, int decode_query_len, float sm_scale, " + "bool linear_softmax, bool use_pdl, bool prewait, Tensor? prefetch_a=None, Tensor? prefetch_b=None, " + "Tensor(d!)? n1_self=None, Tensor? n1_next=None, Tensor? n1_k=None, Tensor? n1_v=None, Tensor? " + "n1_k_scale=None, Tensor? n1_v_scale=None, bool identity=False, bool one_div=False, bool identity_step=False) " + "-> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_sparse_decode_lc8", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SparseDecodeLc8); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp new file mode 100644 index 000000000000..401aa5cc61fd --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void check(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp8_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +} // namespace + +//! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch +//! (kernels/minimaxM3Sm107Tp8/specOffsets.h). +void minimaxM3Sm107Tp8SpecOffsets( + torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) +{ + check(pos, "pos"); + check(posOff, "pos_off"); + check(kv, "kv"); + check(kvOff, "kv_off"); + TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), + "minimax_m3_sm107_tp8_spec_offsets: each tensor and its offsets must have the same length"); + if (pos.numel() == 0 && kv.numel() == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{pos.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSpecOffsets(pos.data_ptr(), + posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), + at::cuda::getCurrentCUDAStream(pos.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_spec_offsets launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SpecOffsets); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp new file mode 100644 index 000000000000..f340b50dc7dd --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp8/stepPrologue.h). +//! +//! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the +//! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses +//! the caller guarantees are live. +//! @param epoch Optional one-element int64 step counter, incremented once. +void minimaxM3Sm107Tp8StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) +{ + TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), + "minimax_m3_sm107_tp8_step_prologue: window must be a contiguous uint8 CUDA tensor"); + TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, + "minimax_m3_sm107_tp8_step_prologue: window is shorter than its header and max_desc descriptors"); + c10::cuda::CUDAGuard const deviceGuard{window.device()}; + int64_t* epochPtr = nullptr; + if (epoch.has_value()) + { + TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 + && epoch->device() == window.device(), + "minimax_m3_sm107_tp8_step_prologue: epoch must be a one-element int64 tensor on the window's device"); + epochPtr = epoch->data_ptr(); + } + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeStepPrologue( + window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_step_prologue launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8StepPrologue); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp new file mode 100644 index 000000000000..8b5e59751f24 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp8/swigluOai.h): x bf16 [rows, 2d] -> out bf16 +//! [rows, d], PDL-launched with an early release when `pdl`. +void minimaxM3Sm107Tp8SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp8_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp8_swiglu_oai: x rows must be 4-byte aligned with an even half width"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp8_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSwigluOai( + reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, + at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_swiglu_oai launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SwigluOai); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py index 05bb489ad81f..421e226b26bf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py @@ -89,6 +89,7 @@ class Arch(str, enum.Enum): SM_100 = "sm_100" SM_103 = "sm_103" + SM_107 = "sm_107" @dataclass(frozen=True) 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 index c47398888520..7576ce1fe445 100644 --- 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 @@ -99,3 +99,85 @@ def _( residual.new_empty(((m + 127) // 128 * 128 * (h // 32),), dtype=torch.uint8), ] return outputs + + +# The MiniMax-M3 sm_107 TP2 / TP4 / TP8 decode ops: one producer and one all-reduce + norm op per TP width (Q and +# index-Q heads per rank from the kernels' producerPw{2,4,8}.h), and the dense QK-norm + RoPE ops. +def _register_minimax_m3_sm107_fakes(tp: int, heads_q: int, heads_index: int) -> None: + @torch.library.register_fake(f"trtllm::minimax_m3_sm107_tp{tp}_producer_pw") + def _( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + prewait_step_inputs=False, + ): + # [Q, index Q] in FP8; the K / V / index-K rows go into the caches in place. + num_tokens = packed.shape[0] + return [ + packed.new_empty((num_tokens, heads_q, 128), dtype=torch.float8_e4m3fn), + packed.new_empty((num_tokens, heads_index, 128), dtype=torch.float8_e4m3fn), + ] + + @torch.library.register_fake(f"trtllm::minimax_m3_sm107_tp{tp}_mc_allreduce_norm") + def _( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + prepushed, + opts, + cap_out=None, + cap_w=None, + cap_eps=0.0, + ): + # [normed, updated residual] (+ [E4M3 rows, swizzled scale bytes] with quant); cap_out is written in place. + outputs = [torch.empty_like(residual), torch.empty_like(residual)] + if quant: + m, h = residual.shape + outputs += [ + residual.new_empty((m, h), dtype=torch.float8_e4m3fn), + residual.new_empty(((m + 127) // 128 * 128 * (h // 32),), dtype=torch.uint8), + ] + return outputs + + +for _tp, _heads_q, _heads_index in ((2, 32, 2), (4, 16, 1), (8, 8, 1)): + _register_minimax_m3_sm107_fakes(_tp, _heads_q, _heads_index) + + +def _minimax_m3_sm107_dense_qk_norm_rope_fp8_fake(qkv, *args, **kwargs): + # A new FP8 tensor shaped like qkv; the KV variant also stores K / V (and the compact Q) in place. + return torch.empty_like(qkv, dtype=torch.float8_e4m3fn) + + +for _name in ( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8", + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv", + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8", + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv", +): + torch.library.register_fake(f"trtllm::{_name}")(_minimax_m3_sm107_dense_qk_norm_rope_fp8_fake) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py new file mode 100644 index 000000000000..25c6fb9961de --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__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_107 / tp2.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/__init__.py new file mode 100644 index 000000000000..3a3bd2995b75 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107 / tp2 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107__tp2/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_107__tp2/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107__tp2/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_107__tp2/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/op.py new file mode 100644 index 000000000000..e0caa5e350dc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/op.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host side of the MiniMax-M3 decode embedding + first-layer RMSNorm (m3_embedding_firstnorm_kernel.py): H 6144, +1..32 rows (one CTA per row: the decode layer's 1..8 rows and the batch-4 / batch-8 plans' 9..32), into caller-owned +buffers. + +Compiled on the first call per (rows, id dtype, table shape, PDL), which must happen outside CUDA-graph capture +(:class:`DecodeEmbedding` compiles at construction). +""" + +import threading +from typing import Dict, Tuple + +import cuda.bindings.driver as cuda_driver +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +MAX_ROWS = 32 # capacity: fixed: the kernel's built rows (decode steps) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def arg(tensor): + return from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _check(ids, table, weight, raw, out) -> None: + rows = ids.shape[0] if ids.dim() == 1 else 0 + if not ( + 1 <= rows <= MAX_ROWS + and ids.dtype in (torch.int32, torch.int64) + and ids.is_contiguous() + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN,) + and all( + t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() + for t in (raw, out) + ) + ): + raise ValueError( + f"m3_embed_norm: expects ids int32/int64 [1..{MAX_ROWS}], a contiguous bf16 table [V, 6144], " + "bf16 weight [6144] and contiguous bf16 raw / out [rows, 6144]" + ) + + +def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) + fn = _compiled[key] = cute.compile( + _m3_embed_norm_jit, + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + _check(ids, table, weight, raw, out) + fn = _compile(ids, table, weight, eps, raw, out, use_pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return out, raw + + +class DecodeEmbedding: + """Caller-owned buffers of one (rows, id dtype) decode embedding: ``raw`` (the gathered rows, the first residual) + and ``output`` (their RMSNorm), over ``table`` and the first layer's plain (non-Gemma) input norm ``weight`` / + ``eps``. Built (and compiled) before CUDA-graph capture.""" + + def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): + assert not torch.cuda.is_current_stream_capturing() + assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN + assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) + assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( + table, + weight, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(eps) + self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) + + def __call__(self, ids): + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return m3_embed_norm( + ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/long_select.py new file mode 100644 index 000000000000..14e24fbdbcac --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/long_select.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. + +Same contract and bit-identical output as ``torch.ops.trtllm.minimax_m3_select_blocks`` on per-KV-head scores +[num_kv_heads, n_blocks, total_q]: init / local forcing, per-query valid-block count, ties to the smaller block index, +forced -inf picks emitted as -1, ascending block ids with -1 tail padding. + +One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups ``block % 16`` is a lower +bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared memory +and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every candidate +against the whole row (re-read). + +The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long`` (this target's copy of ``minimax_m3_select_long`` +with a choice of where it releases the sparse decode), called through this target's ``m3_select_long`` entry, or with +``dsl`` the same selection from one 384-thread CTA per row (``m3_select_long_dsl``, kernels/m3_select_long_dsl). +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long +from ..ops.attention.m3_select_long_dsl import m3_select_long_dsl + + +def long_select_blocks( + max_score_kv: torch.Tensor, + *, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + dsl: bool = False, +) -> Optional[torch.Tensor]: + """Top-16 select over per-KV-head scores [num_kv_heads, n_blocks, total_q] with 129..9216 blocks. + + Returns the head-major backed [total_q, num_kv_heads, 16] int32 table, or None outside this kernel's shapes + (callers keep the reference op). ``early_trigger`` releases the dependent sparse decode right after the select's + grid wait (modeling's ``SELECT_EARLY_MAX_ROWS``); otherwise ``trigger_phase`` 1 releases it once every score of the + row is loaded and 0 at the select's exit (modeling's ``SELECT_LONG_TRIGGER_PHASE``). ``dsl`` runs the 384-thread + kernel (modeling's ``SELECT_LONG_DSL``; the same table). + """ + if ( + max_score_kv.dtype != torch.float32 + or not max_score_kv.is_cuda + or max_score_kv.dim() != 3 + or not 128 < max_score_kv.shape[1] <= MAX_BLOCKS + or n_valid_blocks.dtype != torch.int32 + or not n_valid_blocks.is_cuda + or n_valid_blocks.shape[0] < max_score_kv.shape[2] + ): + return None + num_kv_heads, _, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, 16), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + return (m3_select_long_dsl if dsl else m3_select_long)( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_attn_lc.py new file mode 100644 index 000000000000..ddfbb0e4b2a9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_attn_lc.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 long-context sparse decode attention in one launch. + +Replaces the generic chain's split-K producer (``_gqa_sparse_decode_kernel``, one block per chunk) and the LSE +merge (``_merge_topk_attn_out_kernel``, with its MXFP8 copy for the o_proj) once the top-k block table exists. + +Grid: one cluster of 16 CTAs per (query token, KV head); CTA ``r`` attends the token's ``r``-th selected block of +that head (the Triton producer's chunk ``r`` at 16 chunks), keeps the partial in shared memory, and after a cluster +barrier merges head ``r`` of the KV head's group from all 16 partials over DSMEM, writing the BF16 row and, +optionally, its MXFP8 copy and swizzled UE8M0 scales, bitwise ``torch.ops.trtllm.mxfp8_quantize(out, True)``. + +Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +log-domain recurrence. Accumulation order inside the MMAs and the reductions differs, so results match to rounding, +not bit for bit. + +Schedule: the request length and the query rows (FP8 widened to FP16 in shared memory) are read before the PDL grid +wait, which overlaps them with the select ahead when it releases this grid early (modeling's ``SELECT_EARLY_MAX_ROWS``); +the split-K partials travel by ``st.async`` into each receiving CTA's own mbarrier instead of a cluster-wide barrier. +Neither changes a value or the merge order. + +Geometry: 16 query heads per KV head, two KV heads per rank (TP2), head_dim 128, page 128, top-k 16. The kernel +is ``torch.ops.trtllm.minimax_m3_sm107_tp2_sparse_decode_lc``, called through this target's ``m3_sparse_decode_lc`` +entry. +""" + +import torch + +from tensorrt_llm.logger import logger + +from ..ops.attention.m3_sparse_decode_lc import m3_sparse_decode_lc + + +class N1Slot: + """The sparse layer's N1 id buffer (modeling's ``SPARSE_DECODE_LC_PREFETCH_SELF_KV``), set by the layer around its + decode call and taken by the long-context decode; empty means no prefetch.""" + + def __init__(self) -> None: + self.ids = None + + def take(self): + ids, self.ids = self.ids, None + return ids + + +# int32 [2, N1_MAX_ROWS, 16]: the layer's selected block ids per (KV head, query row, slot) at its previous decode step. +N1_IDS = N1Slot() + + +def lc_applies(q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output) -> bool: + """Whether the step's operands fit the kernel: TP2 geometry, packed NVFP4 pages, FP8 q, BF16 [T, 32, 128] out.""" + if q.dim() != 3 or tuple(q.shape[1:]) != (32, 128) or q.dtype != torch.float8_e4m3fn: + return False + if q.stride(2) != 1 or q.stride(1) % 16 or q.stride(0) % 16 or q.data_ptr() % 16: + return False + if ( + int(k_paged.shape[1]) != 2 + or topk_idx.dim() != 3 + or int(topk_idx.shape[0]) != 2 + or int(topk_idx.shape[1]) != int(q.shape[0]) + or int(topk_idx.shape[-1]) != 16 + or topk_idx.dtype != torch.int32 + or topk_idx.stride(2) != 1 + ): + return False + if v_paged.stride() != k_paged.stride(): + return False + for t in (k_paged, v_paged): + if ( + t.stride(3) != 1 + or t.stride(2) != 64 + or int(t.shape[3]) != 64 + or t.data_ptr() % 16 + or t.stride(0) % 16 + or t.stride(1) % 16 + ): + return False + for t in (k_block_scale, v_block_scale): + if ( + t is None + or int(t.shape[1]) != 2 + or t.stride(3) != 1 + or t.stride(2) != 8 + or t.data_ptr() % 8 + or t.stride(0) % 8 + or t.stride(1) % 8 + ): + return False + if v_block_scale.stride() != k_block_scale.stride(): + return False + return ( + output.dtype == torch.bfloat16 + and output.dim() == 3 + and output.is_contiguous() + and tuple(output.shape) == tuple(q.shape) + ) + + +def m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + linear_softmax=False, + mxfp8=None, + use_pdl=True, + prewait=False, + n1_ids=None, +): + """Drop-in for ``minimax_m3_sparse_attn_decode`` on the TP2 NVFP4 shape; callers check ``lc_applies`` first. + + ``topk_idx`` is the head-major [2, total_q, 16] table the Triton producer takes. ``mxfp8`` is ``(q8, sf)`` from + ``MERGE_MXFP8.buffers(4096, dev)`` to also emit the o_proj hand-off. ``n1_ids`` is the layer's N1 id buffer + (``N1_IDS``): the decode prefetches last step's blocks into L2 and records this step's. + """ + total_q = int(q.shape[0]) + q8, sf = ( + (mxfp8[0].view(torch.uint8), mxfp8[1].view(torch.uint8)) + if mxfp8 is not None + else (None, None) + ) + m3_sparse_decode_lc( + q.view(torch.uint8), + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + output.view(total_q, 4096), + q8, + sf, + 4096 // 128, + int(decode_query_len), + float(sm_scale), + bool(linear_softmax), + bool(use_pdl), + bool(prewait), + n1_ids, + ) + key = (total_q, q8 is not None, n1_ids is not None) + if key not in _engaged: + _engaged.add(key) + logger.info( + f"[M3 sm_107 tp2] sparse decode Lc engaged: T={total_q}, kv_heads=2, mxfp8_handoff={q8 is not None}, " + f"n1_self_kv={n1_ids is not None}" + ) + + +_engaged = set() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_draft_step.py new file mode 100644 index 000000000000..c61ab4acbf14 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_draft_step.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 draft pass for the MiniMax-M3 drafter's vocab-sharded BF16 head (TP decode, <= 8 requests). + +One kernel replaces a draft pass's head and the draft loop's per-pass bookkeeping, at the in-tree precision: + +* The drafter's BF16 lm_head shard is kept a second time, built by ``prepare_head`` before the executor's warmup: in + the lossless 12-bit coding of the target head (``ops.gemm.m3_target_head``; 75% of the bytes, kept only when it + decodes to the BF16 shard bit for bit), else in ``mma.m16n8k16`` A-fragment order (``bf16_fragments``, the same + values). The kernel streams it with BF16 MMAs and FP32 accumulation over the BF16 head input rows; a coded + exception decodes to +0 and its exact product x * w is added at its tile's flush. +* Each logit is rounded to BF16 (the in-tree logits are the BF16 head GEMM's output) and keyed + ``ordered(value) << 32 | ~global_index`` with -0 as +0, so the larger value wins and ties take the lowest index, as + the in-tree ``torch.max`` / ``torch.argmax`` over the gathered logits (``greedy_sample_draft_with_tp_gather``). +* The last CTA exchanges each row's key across the TP ranks through a multicast mapping (Lamport slots, see + ``m3_shard_accept``), then stores the draft token, the next pass's input id and position, and advances the draft KV + lengths as the in-tree draft loop does between passes. + +The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu) and called through the +catalog entry ``ops.spec.m3_draft_step``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok +from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments +from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op +from . import m3_report +from .m3_shard_accept import AcceptWorkspace, create_workspace + +_ENGAGED: set = set() + + +# The fused draft loop streams the 12-bit coded copy of the draft head when it decodes bit for bit; off at TP2 (the +# coded head is slower on sm_107: the BF16 fragments stream). +HEAD_CODED = False +# The head copy is placed per locality domain and each pass runs as two launches, one per domain (placement only; +# 105.6 -> 86.3 / 86.6 us per pass at b1 / b2 on Rubin TP2). +UGPU = True +# The head streams an equal share of each launch's (tile, chunk) space per CTA instead of whole 16-row tiles (two +# zero-filled addends per boundary tile, so the sum does not depend on their order). +EQ_CHUNKS = True + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the 12-bit coded stream + with its exception lists (``coded``), or the fragment-ordered BF16 copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + coded: bool = False + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + exp_base: int = 0 + # BF16 fragments placed per locality domain (UGPU below): tiles [0, split_tile) live in domain 0 and + # the rest in domain 1; 0 = not placed (one launch over all SMs). + split_tile: int = 0 + # Equal-chunk split (EQ_CHUNKS below): zero-filled shared-tile partials [tiles, 32, 4] fp32 and + # arrival counts [tiles] int32 (the kernel leaves them zero); None = whole tiles per CTA. + bnd: Optional[torch.Tensor] = None + bnd_cnt: Optional[torch.Tensor] = None + + +def head_supported(lm_head) -> bool: + """A bias-free BF16 head whose local shard the kernel can own: whole 16-row tiles, K a multiple of 16, at most + MAX_TILES_PER_CTA tiles per CTA of a one-CTA-per-SM grid.""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + if getattr(lm_head, "bias", None) is not None: + return False + n, k = weight.shape + if n % 16 or k % 16: + return False + return (n // 16 + _num_sms() - 1) // _num_sms() <= MAX_TILES_PER_CTA + + +def prepare_head(lm_head, tp_rank: int) -> DraftHead: + """The streamed copy of ``lm_head``'s local shard (global ids from tp_rank * rows): its 12-bit coding when + ``HEAD_CODED``, K is a multiple of 64 and the coding decodes to the shard bit for bit, else (fail + closed) the fragment-ordered BF16 copy. + Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_draft_head``).""" + weight = lm_head.weight + offset = tp_rank * weight.shape[0] + with torch.no_grad(): + w = weight.detach() + if HEAD_CODED and w.shape[1] % 64 == 0: + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return DraftHead( + weight=weight, + wf=packed, + vocab_offset=offset, + coded=True, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + exp_base=base, + bnd=None, # the coded head (off on sm_107) keeps whole tiles + bnd_cnt=None, + ) + del packed, exc_off, exc_pos, exc_w + wf = bf16_fragments(w) + wf, split_tile = _place_per_domain(wf) + bnd, bnd_cnt = _eq_buffers(wf.shape[0], wf.device) + return DraftHead( + weight=weight, wf=wf, vocab_offset=offset, split_tile=split_tile, bnd=bnd, bnd_cnt=bnd_cnt + ) + + +def _eq_buffers(n_tiles: int, device): + """The equal-chunk split's zero-filled shared-tile partials and arrival counts, or (None, None) when + EQ_CHUNKS is off.""" + if not EQ_CHUNKS: + return None, None + return ( + torch.zeros(n_tiles, 32, 4, dtype=torch.float32, device=device), + torch.zeros(n_tiles, dtype=torch.int32, device=device), + ) + + +def _place_per_domain(wf: torch.Tensor): + """(wf, split tile): the BF16 fragments split-mapped over the two locality domains at the granule-aligned tile + nearest half (locality.LocalityMap.place_split; its domain streams built here, before capture), or (wf, 0) when + UGPU is off. Placement only: the streamed values are the same.""" + if not UGPU: + return wf, 0 + from .m3_locality import LocalityMap + + lm = LocalityMap.get(wf.device) + tile_bytes = wf[0].numel() * wf.element_size() + step = 1 + while (step * tile_bytes) % lm.granularity: + step += 1 + n_tiles = wf.shape[0] + split_tile = max(step, round(n_tiles / 2 / step) * step) + per = lm.sms_per_domain + if split_tile >= n_tiles or max(split_tile, n_tiles - split_tile) > per * MAX_TILES_PER_CTA: + return wf, 0 + lm.domain_streams() + return lm.place_split(wf, split_tile), split_tile + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace ([2][MAX_TOKENS][world] slots), created collectively on the first eager call; a + CUDA-graph capture that finds it missing is a miss.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("m3_draft_step", None, "the draft-step exchange workspace") + return None + _WS = create_workspace(mapping, device, MAX_TOKENS) + return _WS + + +def draft_step( + head: DraftHead, + x: torch.Tensor, + rows: int, + x_rows: Optional[torch.Tensor], + mapping, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, +) -> bool: + """One fused greedy draft pass over ``rows`` requests (see the module docstring); False when the exchange + workspace is missing during a capture. Collective over the TP group. + + Args: + head: the prepared head of this rank (``prepare_head``), coded or BF16 fragments. + x: [*, K] BF16 head input rows (the drafter's final-norm output); row ``x_rows[b]`` (int64) or ``b`` is + request b's. + draft_out: int32 [rows, >= step + 1], unit column stride: column ``step`` receives the tokens. + ids_next / pos_next: int32 [rows]: the tokens, and ``pos_src[x_rows[b] or b] + 1`` (pos_src int32). + kv_lens: int32 draft KV lengths: ``-= draft_len - num_acc`` with ``seq_lens = 1`` on pass 0 when ``num_acc`` + is given, ``+= 1`` on later passes. + """ + ws = workspace(mapping, x.device) + if ws is None: + return False + if head.split_tile: + _draft_step_split( + head, + x, + rows, + x_rows, + ws, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + ) + return True + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + _num_sms(), + True, + None, + head.exc_off, + head.exc_pos, + head.exc_w, + head.exp_base, + 0, + 0, + 0, + head.bnd, + head.bnd_cnt, + ) + return True + + +def _draft_step_split( + head, + x, + rows, + x_rows, + ws, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, +) -> None: + """One pass as two launches, one per locality domain on that domain's green-context stream, each streaming its own + domain's tiles of the split-mapped fragments; the keys' max and the last-CTA exchange and bookkeeping run once over + both (the last of the 2 x SMs-per-domain CTAs). Forked from and joined back to the current stream; no PDL on the + forks.""" + from .m3_locality import LocalityMap + + lm = LocalityMap.get(x.device) + per = lm.sms_per_domain + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + streams = lm.domain_streams() + n_tiles = head.wf.shape[0] + if x.device not in _ENGAGED: + _ENGAGED.add(x.device) + logger.info( + f"[M3 sm_107 tp2] fast path engaged: m3_draft_step_ugpu (split tile {head.split_tile} of {n_tiles}, " + f"{per} CTAs per domain, on {x.device})" + ) + for st, lo, hi in ((streams[0], 0, head.split_tile), (streams[1], head.split_tile, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + per, + False, + None, + None, + None, + None, + 0, + lo, + hi, + 2 * per, + head.bnd, + head.bnd_cnt, + ) + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_drafter_attn.py new file mode 100644 index 000000000000..eabe675e6970 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_drafter_attn.py @@ -0,0 +1,413 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's decode attention on Rubin: ``M3DrafterFullCtxFmha``. + +The drafter layer is a plain ``TrtllmAttention`` over its own FP8 paged cache (the full context, every draft pass). +On SM 10.7 its generation phase runs ``thop.attention``'s in-tree trtllm-gen runner, whose split is capped at 8 CTAs +per (request, KV head): 45-46 us per pass at 165k tokens on 212 SMs, for 85 MB. + +This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append +into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context +PrimTS split-KV decode of ``kernels/m3_fullctx_decode`` at the same pinned recipes as the dense layers. + +With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead +(``minimax_m3_sm107tp2_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page +pool and block table taken from the KV cache pool pointers and mapping the preprocess reads. + +It serves generation-only steps of the covered (query rows per request, requests) shapes; the target installs it +first in the drafter layer's FMHA list (``install``), so every other step falls through to the stock libraries. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + FlashInferTrtllmGenFmha, +) +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode + +from . import m3_fullctx_decode as fullctx + +# The fused RoPE + FP8 K/V launch also takes a one-request multi-token pass (the drafter's first pass at one +# request), where the in-tree preprocess needs the token-map guard: bitwise the guarded preprocess + cast. +ROPE_ONE_REQUEST_MULTI = True +# With PDL, that launch takes the maximum shared-memory carveout so the drafter decode's large-smem CTAs co-reside with +# it (timing only). +ROPE_CARVEOUT = True + + +def fwd_has_no_mrope(fwd) -> bool: + return getattr(fwd, "mrope_position_deltas", None) is None + + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + +class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): + """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + + def __init__( + self, + attn: "TrtllmAttention", + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + pdl: bool = False, + on_rope_taken=None, + ) -> None: + super().__init__(attn) + self._recipes = recipes + self._max_kv_len = max_kv_len + self._own_reducer = own_reducer + self._fused_rope_mode = fused_rope_mode + self._pdl = pdl + self._on_rope_taken = on_rope_taken + self._one_request_multi = ROPE_ONE_REQUEST_MULTI + self._on_taken = on_taken + self._on_uncovered = on_uncovered + self._on_missed = on_missed + + def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: + nreq = int(metadata.num_generations) + if nreq <= 0 or int(q.shape[0]) % nreq: + return None + return int(q.shape[0]) // nreq, nreq + + def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: + if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: + return False + if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: + return False + attn = self.attn + if ( + attn.num_heads != fullctx.HEADS + or attn.num_kv_heads != fullctx.KV_HEADS + or attn.head_dim != fullctx.HEAD_DIM + or not QuantMode(attn.quant_mode).has_fp8_kv_cache() + ): + self._on_uncovered("not the TP2 FP8 drafter geometry") + return False + # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's + # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. + if ( + forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant + or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig + ): + self._on_uncovered("per-step KV cache scales") + return False + if int(metadata.max_seq_len) > self._max_kv_len: + self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") + return False + shape = self._shape(q, metadata) + if shape is None or shape not in self._recipes: + self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") + return False + # Compiled before warmup (``install``); the forward never compiles one. + if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: + self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") + return False + return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) + + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + attn = params.attn + width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim + rotary_dim = int(attn.rope_params.dim) + return ( + self._fused_rope_mode is not None + and params.qkv_input is not None + and attention_input.dim() == 2 + and attention_input.dtype == torch.bfloat16 + and attention_input.is_contiguous() + and int(attention_input.shape[1]) == width + and not params.is_cross + and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) + and attn.rotary_cos_sin is not None + and attn.rotary_cos_sin.dtype == torch.float32 + and fwd_has_no_mrope(params.fwd) + and 0 < rotary_dim <= attn.head_dim + and rotary_dim % 16 == 0 + and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) + # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo + # runs for several requests only) and reads the workspace's; with ROPE_ONE_REQUEST_MULTI + # the fused launch (which places rows by request = token / q_len, no map) takes it too, else it keeps + # the in-tree path (and its token-map guard). + and ( + self._one_request_multi + or not (self._multi_token(params) and int(params.batch_size) == 1) + ) + ) + + @staticmethod + def _multi_token(params: FmhaParams) -> bool: + return ( + params.spec_decoding_generation_lengths is not None + and int(params.attn.predicted_tokens_per_seq) > 1 + ) + + def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" + if not self._fused_rope_applies(params, attention_input): + return False + attn = params.attn + meta = params.meta + qlen, nreq = int(params.input_seq_length), int(params.batch_size) + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + if plan is None: + return False + # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary + # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the + # block table is the pool's [requests, 2, blocks] K / V sub-page indices. + pool_index, layer_in_pool = ( + int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] + ) + tpb = int(params.tokens_per_block) + sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim + base = ( + int(meta.host_kv_cache_pool_pointers[pool_index, 0]) + + layer_in_pool * int(params.kv_factor) * sub_page_bytes + ) + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) + ) + tables = meta.kv_cache_block_offsets[ + pool_index, params.seq_offset : params.seq_offset + nreq + ] + multi = self._multi_token(params) + if not torch.cuda.is_current_stream_capturing(): + # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess + # path checks; the op itself writes V at the table's V sub-page. + tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(tables.shape[-1], device=tables.device) + in_use = cols[None, :] < tpb_held[:, None] + offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + q = torch.empty( + int(params.num_tokens), + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=attention_input.device, + ) + torch.ops.trtllm.minimax_m3_sm107tp2_drafter_rope_fp8_kv( + attention_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + int(attn.rope_params.dim), + tpb, + self._fused_rope_mode, + self._pdl, + ROPE_CARVEOUT, + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + tables[:, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + if self._on_rope_taken is not None: + self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") + self._on_taken(params.num_tokens) + return True + + def run_generation(self, params: FmhaParams) -> None: + attention_input = params.qkv_input if params.qkv_input is not None else params.query_input + if attention_input is None: + raise RuntimeError("the drafter decode requires QKV or query input.") + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + if self._run_fused_rope(params, attention_input): + return + # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 + # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). + fp8_context_fmha = False + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + _fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + _window_left, + is_multi_token_gen, + ) = thop.trtllm_gen_generation_preprocess( + attention_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + fp8_context_fmha, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + ) + qlen, nreq = params.input_seq_length, params.batch_size + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the + # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row + # pass is never multi-token. + if plan is None or (is_multi_token_gen and qlen == 1): + raise RuntimeError( + f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" + ) + rows = params.num_tokens * attn.num_heads * attn.head_dim + q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) + q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) + # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 + # of the pool (the two views have the same page count, as the kernel requires). + if not torch.cuda.is_current_stream_capturing(): + # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. + lengths = params.spec_decoding_generation_lengths + if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): + raise RuntimeError( + f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" + ) + # Only the blocks a request holds are laid out; the table's tail entries are unset. + tpb = int(params.tokens_per_block) + held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(block_tables.shape[-1], device=block_tables.device) + in_use = cols[None, :] < held[:, None] + offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. + # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + block_tables[:nreq, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + self._on_taken(params.num_tokens) + + +def install( + draft_model: torch.nn.Module, + recipes, + max_kv_len: int, + own_reducer: bool, + pdl: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + on_rope_taken=None, +) -> int: + """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention + backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it + was installed on.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + count = 0 + device = next(draft_model.parameters()).device + # The backend is not a module: each attention module of the drafter holds it as ``attn``. + for module in draft_model.modules(): + backend = getattr(module, "attn", None) + if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): + continue + if not bool((backend.kv_cache_scaling_factor == 1.0).all()): + on_uncovered("a scaled FP8 KV cache") + continue + fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) + libs = backend._fmha_manager.fmha_libs + if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): + libs.insert( + 0, + M3DrafterFullCtxFmha( + backend, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode=fused_rope_mode, + pdl=pdl, + on_rope_taken=on_rope_taken, + ), + ) + count += 1 + return count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fold_tail.py new file mode 100644 index 000000000000..1736115f1386 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fold_tail.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The speculative step's store and next-step staging at the tail of the forward (one Triton launch). + +For each of the first ``n`` rows (``n`` read from device memory, so a captured graph whose step was not armed writes +nothing), with ``s = slots[i]``: + +* the spec sampler's per-slot store, exactly as ``m3_spec_tail.slot_scatter`` writes it (zero-padded to each store + width): ``new[s][c]``, ``next_new[s][c]``, ``draft[s][c]`` and ``lens[s]``; +* the engine's next-step inputs, exactly as ``m3_spec_tail.step_gather`` writes them from that store for a step whose + rows are these rows in this order (no new tokens, first previous row 0): ``ids[i*T + t] = next_new[s][t]``, + ``draft_ids[i*D + d] = draft[s][d]``, ``pos_off[i*T + t] = lens[s]``, ``kv_off[i] = lens[s] - T``. + +Integer copies only: the values are the ones the sampler's and the engine's own copies write. + +With ``cpp`` the same writes run as this target's C++ op (trtllm::minimax_m3_sm107_tp2_fold_tail, +minimaxM3Sm107Tp2/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fold_tail_kernel( + ctrl, + slots, + src_new, + sn_s0, + sn_s1, + SN_W: tl.constexpr, + dst_new, + dn_s0, + dn_s1, + DN_W: tl.constexpr, + src_next, + sx_s0, + sx_s1, + SX_W: tl.constexpr, + dst_next, + dx_s0, + dx_s1, + DX_W: tl.constexpr, + src_draft, + sd_s0, + sd_s1, + SD_W: tl.constexpr, + dst_draft, + dd_s0, + dd_s1, + DD_W: tl.constexpr, + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + T: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(ctrl) + if i < n: + s = tl.load(slots + i).to(tl.int64) + c = tl.arange(0, BLOCK) + new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) + tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) + nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) + tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) + drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) + tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) + ln = tl.load(src_lens + i) + tl.store(dst_lens + s, ln) + # The next step's inputs read the store's padded rows: zero past each source width. + tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) + tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) + tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) + tl.store(kv_off + i, ln - T) + + +def fold_tail( + ctrl: torch.Tensor, + slots: torch.Tensor, + rows: int, + src_new: torch.Tensor, + dst_new: torch.Tensor, + src_next: torch.Tensor, + dst_next: torch.Tensor, + src_draft: torch.Tensor, + dst_draft: torch.Tensor, + src_lens: torch.Tensor, + dst_lens: torch.Tensor, + ids: torch.Tensor, + draft_ids: torch.Tensor, + pos_off: torch.Tensor, + kv_off: torch.Tensor, + tokens_per_step: int, + cpp: bool = False, +) -> bool: + """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs + are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width].""" + two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) + one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) + if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): + return False + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) + or slots.dtype != torch.long + ): + return False + t_step = tokens_per_step + d_step = t_step - 1 + if ( + rows <= 0 + or slots.numel() < rows + or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) + or ids.numel() < rows * t_step + or pos_off.numel() < rows * t_step + or draft_ids.numel() < rows * d_step + or kv_off.numel() < rows + ): + return False + widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) + widths += (src_draft.shape[1], dst_draft.shape[1], t_step) + if cpp and max(widths) <= 32: + torch.ops.trtllm.minimax_m3_sm107_tp2_fold_tail( + ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, + src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, + ) # fmt: skip + return True + block = triton.next_power_of_2(max(widths)) + _fold_tail_kernel[(rows,)]( + ctrl, + slots, + src_new, + src_new.stride(0), + src_new.stride(1), + src_new.shape[1], + dst_new, + dst_new.stride(0), + dst_new.stride(1), + dst_new.shape[1], + src_next, + src_next.stride(0), + src_next.stride(1), + src_next.shape[1], + dst_next, + dst_next.stride(0), + dst_next.stride(1), + dst_next.shape[1], + src_draft, + src_draft.stride(0), + src_draft.stride(1), + src_draft.shape[1], + dst_draft, + dst_draft.stride(0), + dst_draft.stride(1), + dst_draft.shape[1], + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + t_step, + d_step, + block, + num_warps=1, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_force_vec.py new file mode 100644 index 000000000000..cc0ce28d9639 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_force_vec.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's +``force_num_accepted_tokens``), in one launch. Reached through the ``ops/spec`` entries ``m3_force_vec`` and +``m3_force_vec_req``. + +Triton. ``force_vec_kernel``: one CTA bumps the shared step counter and draws each request's count from the fixed-seed +pool with the in-tree (counter, slot) hash, exactly as ``_apply_force_accepted_tokens`` does. ``force_vec_req_kernel``: +one draw per request from its own key (its KV length at this verify), no shared counter. +""" + +import triton +import triton.language as tl + + +# One compiled variant for every batch and accept count: specializing them (== 1, divisible by 16) would compile +# mid-serving the first time a step reaches a new one, past the worker's warm-up. +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_kernel( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK: tl.constexpr, +): + c = tl.load(counter) + 1 + slots = tl.arange(0, BLOCK) + tl.store(counter, c) + idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask + live = slots < batch_size + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) + + +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_req_kernel( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK: tl.constexpr, +): + # One draw per request from its own key (its KV length at this verify): no shared counter, so a request's forced + # counts depend only on its own prompt and progress, not on what else the server ran before or beside it. + slots = tl.arange(0, BLOCK) + live = slots < batch_size + k = tl.load(keys + slots, mask=live, other=0).to(tl.int64) + idx = (k * key_stride) & pool_mask + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fullctx_decode.py new file mode 100644 index 000000000000..c220741e3874 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fullctx_decode.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 full-context FP8 decode attention on Rubin: ``fullctx_decode``. + +The dense layers (0-2) attend every token of their FP8 E4M3 (unscaled) paged KV cache. On SM 10.7 the stock call +(flashinfer trtllm-gen, ``minimax_m3_trtllm_gen_dense_decode``) caps its split at 16 cluster CTAs per (request, KV +head): at TP2 (2 KV heads) one request runs on 32 CTAs of 212 and takes 28-30 us for 85 MB at 165k tokens. + +This runs the same arithmetic -- FP8 Q (the same E4M3 rounding of q), FP8 K/V, P quantized to E4M3 per split -- on +the vendored PrimTS decode kernel (``attention/backends/prims_ts``, SM 10.7 through the sm_100f family) at a pinned +split-KV recipe, which fans one request out over 96-144 CTAs. The split partials are merged by +``torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce`` (``cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/ +fullCtxReduce.cu``): unnormalized fp32 accumulation over the active splits, one final division, one BF16 rounding, +the merge form of stock trtllm-gen. Only the split changes against stock (modeling's ``FULLCTX_DECODE_RECIPES``). + +A plan is compiled per (device, query rows per request, requests) by ``prebuild`` before the executor's warmup; the +forward only reads plans (``plan_for``). The plan is built from PrimTS's own config, compile and workspace helpers at +an explicit recipe; the PrimTS wrapper's planning policy is not used, so no PrimTS state is replaced. With +``own_reducer`` off the PrimTS reducer merges into FP16 and the result is cast to BF16. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Mapping, Optional, Tuple + +import torch + +HEADS = 32 +KV_HEADS = 2 # capacity: fixed: the tp2 KV heads per rank +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim +PAGE = 128 +# The per-request K/V length PrimTS's config selection sizes its policy for (Q/KV tiles, launch mode); the pinned +# recipes were measured with it. The compiled capacity (``max_kv_len``) may be larger. +POLICY_KV_LEN = ( + 1 << 20 +) # capacity: fixed: the PrimTS policy's K/V length bound (the plans' compile shape) +HEADS_PER_KV = HEADS // KV_HEADS # capacity: fixed: the tp2 query heads per KV head + +_plans: Dict[Tuple[int, int, int, bool], "FullCtxPlan"] = {} +_lock = threading.Lock() + + +class FullCtxPlan: + """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" + + def __init__( + self, + device: torch.device, + qlen: int, + nreq: int, + max_kv_len: int, + recipe: Mapping, + own_reducer: bool, + pdl: bool, + ): + from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec + from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( + get_max_active_clusters_for_cluster_size, + make_decode_config, + ) + + self._tsdec = tsdec + self.qlen = qlen + self.nreq = nreq + self.pdl = pdl + self.recipe = dict(recipe) + splits = int(self.recipe["splits"]) + fp8 = tsdec._cutlass_dtype("float8_e4m3fn") + # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer + # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents + # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an + # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the + # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own + # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge + # (minimax_m3_sm107_tp2_fullctx_reduce) waits as its first instruction. + cfg = make_decode_config( + headdim=HEAD_DIM, + args={"use_pdl": pdl, **self.recipe.get("config", {})}, + seq_len_q=qlen, + seq_len_kv=min(max_kv_len, POLICY_KV_LEN), + batch_size=nreq, + num_heads_q=HEADS, + num_heads_kv=KV_HEADS, + q_dtype=fp8, + k_dtype=fp8, + v_dtype=fp8, + o_dtype=tsdec._cutlass_dtype("float16"), + qkv_layout="pagedKv", + num_tokens_per_page=PAGE, + storage_tokens_per_page=PAGE, + split_kv_mode="gmem_reduction_with_separate_kernel", + splits_kv=splits, + max_splits_kv=splits, + min_loop_iters_per_split=1, + mask_type="causal" if qlen > 1 else "dense", + auto_tuner=True, + split_kv=True, + ) + _check_layout(cfg, qlen, splits) + with torch.cuda.device(device): + max_active = get_max_active_clusters_for_cluster_size(1) + spec = tsdec._decode_launch_spec_from_config( + cfg, + batch_size=nreq, + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + seq_len_q=qlen, + max_active_clusters=int(max_active), + ) + compile_spec = tsdec._make_decode_compile_spec( + spec, + device_index=_index(device), + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + page_size=PAGE, + max_kv_len=max_kv_len, + seq_len_q=qlen, + q_dtype_key="float8_e4m3fn", + k_dtype_key="float8_e4m3fn", + v_dtype_key="float8_e4m3fn", + output_dtype_key="float16", + use_packed_q=False, + kv_prefix_mode="dynamic", + kv_lengths_mode="dynamic", + ) + self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) + layout = tsdec._make_decode_workspace_layout( + spec.scratch_shapes, + torch.float16, + use_separate_reduction_kernel=True, + use_split_kv=True, + ) + self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) + self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) + self.workspace.split_kv_counter.zero_() + self.workspace.cu_seqlens_q.zero_() + self.workspace.attention_sinks.zero_() + rows = HEADS_PER_KV * qlen + if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): + raise RuntimeError( + f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " + f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" + ) + self.own_reducer = own_reducer + self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) + self.tile_kv = int(cfg.tile_size_kv) + self.insts_kv = int(cfg.num_insts_kv) + self.causal = qlen > 1 + shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) + self.staging = torch.empty(shape, dtype=torch.float16, device=device) + + def run( + self, + q: torch.Tensor, + k_pages: torch.Tensor, + v_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + ) -> None: + # The kernel takes a compact query; a column view of the fused QKV projection is copied once. + q = q.reshape(self.staging.shape).contiguous() + runtime = self._tsdec._prepare_decode_runtime_unchecked( + q, + (k_pages, v_pages), + None, + output_dtype=torch.float16, + bmm1_scale=sm_scale, + bmm2_scale=1.0, + out=self.staging, + ) + ws = self.workspace + self._main( + runtime.q, + runtime.k_cache, + runtime.v_cache, + runtime.k_sf_cache, + runtime.v_sf_cache, + runtime.out, + seq_lens, + ws.cu_seqlens_q, + block_table, + block_table, # Dense decode constexpr-elides the membership table. + ws.partial_o, + ws.partial_stats, + ws.split_kv_counter, + ws.attention_sinks, + runtime.num_physical_pages, + runtime.k_page_stride, + runtime.k_head_stride, + runtime.k_token_stride, + runtime.v_page_stride, + runtime.v_head_stride, + runtime.v_token_stride, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + if self.own_reducer: + torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce( + ws.partial_o, + ws.partial_stats, + seq_lens, + output, + self.qlen, + self.q_tokens_per_cta, + self.tile_kv, + self.insts_kv, + self.causal, + self.pdl, + ) + return + self._primts_reducer( + runtime.out, + seq_lens, + ws.cu_seqlens_q, + ws.partial_o, + ws.partial_stats, + ws.attention_sinks, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + output.view(self.staging.shape).copy_(self.staging) + + +def _check_layout(cfg, qlen: int, splits: int) -> None: + """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. + + Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is + ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding + window, attention sinks or packed Q. + """ + problems = [] + if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): + problems.append("not split-KV with the separate GMEM reducer") + if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: + problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") + if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: + problems.append("cluster-SMEM reduction or a persistent scheduler") + if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: + problems.append("packed Q, a sliding window or attention sinks") + if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): + problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") + if qlen > 1: + if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): + problems.append("multi-token Q without the grouped [token, head] tile layout") + elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): + problems.append( + f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" + ) + if problems: + raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + + +def has_own_reducer() -> bool: + """Whether this build carries the BF16 split-KV merge op.""" + return hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_fullctx_reduce") + + +def plan_for( + device: torch.device, qlen: int, nreq: int, own_reducer: bool = True +) -> Optional[FullCtxPlan]: + """The prebuilt plan for this shape, or None.""" + return _plans.get((_index(device), qlen, nreq, own_reducer)) + + +def prebuild( + device: torch.device, + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + pdl: bool, +) -> int: + """Compile the plan of every (query rows per request, requests) in ``recipes``. Call before warmup.""" + built = 0 + for (qlen, nreq), recipe in recipes.items(): + key = (_index(device), qlen, nreq, own_reducer) + with _lock: + if key not in _plans: + _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) + built += 1 + return built + + +def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: + """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" + return pool[0::subpages_per_slot], pool[1::subpages_per_slot] + + +def _index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_index_score.py new file mode 100644 index 000000000000..8088d79f014f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_index_score.py @@ -0,0 +1,670 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with four schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. +* ``prestage_old`` (with ``early_issue``, ``INDEX_SCORE_PRESTAGE_OLD``): when the CTA's first block holds only + tokens older than this step (block < (seq_lens hint - decode query length) / 128), its K tile is TMA-loaded into its + stage before the wait; Q (and a first block that is not all old) is still loaded after it. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer``) through ``IndexScoreRunner`` +and called through this target's ``m3_index_score`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, griddepcontrol_wait + +# With the early issue, a CTA whose first block holds only tokens older than this step also TMA-loads that block's +# index-K tile before the grid wait (bit-identical scores; replay period -47 us per step on Rubin TP2). +INDEX_SCORE_PRESTAGE_OLD = True + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + prestage_old: bool = False, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K + # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's + # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. + self.prestage_old = prestage_old and early_issue + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # Otherwise the dependent block selector is released as the CTAs exit. + early_q = False + early_k = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + seqlen_hint = seq_lens[batch_id] + hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + if cutlass.const_expr(self.prestage_old): + # Stage the first K tile now when every token of the block predates this step. + if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: + early_k = True + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + if not early_k: + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # The length that bounds this grid's work is read after the wait (the + # pre-wait reads above are prefetch / early-issue hints only, clamped to + # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens + # patches every step, and num_blocks bounds an unmasked block_table load + # and an unmasked score store, so a stale length is an out-of-bounds write + # rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + prestage_old=early and INDEX_SCORE_PRESTAGE_OLD, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_locality.py new file mode 100644 index 000000000000..3e88c7fb1ce2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_locality.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Locality domains ("uGPU") of a multi-die GPU (Rubin VR200: 2 domains of 100 SMs + 12 remainder SMs) for the +streaming decode kernels of this target: the SM -> domain map, per-domain weight placement, and the CTA -> SM +assignment a persistent kernel uses to read only its own die's memory. + +A kernel launched with one CTA per SM reads ``%smid`` once and takes its logical block index from a host-built +permutation (:meth:`LocalityMap.logical_map`), so every per-CTA structure (work plan row, epochs, roles) follows its +die. Weights that are split by output rows into per-domain halves (:meth:`LocalityMap.split_rows`) live in one virtual +range whose halves are physically in the two domains. Placement only: the arithmetic is unchanged, so results are +bitwise those of the unplaced weights. + +The SM -> domain map is attested by the driver: the SMs are split by locality domain +(``cuDevSmResourceSplit`` with ``CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID``), and a ``%smid`` probe runs on a green +context of each split. Allocations use driver VMM with ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN`` and are checked +with ``CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL``; any failure raises (fail closed).""" + +from __future__ import annotations + +import threading +from typing import Dict, List, Sequence, Tuple + +import torch + +_LOCK = threading.Lock() +_MAPS: Dict[int, "LocalityMap"] = {} +_PROBE_SRC = r""" +extern "C" __global__ void smid_probe(int* out) { + if (threadIdx.x == 0) { + unsigned s; + asm volatile("mov.u32 %0, %%smid;" : "=r"(s)); + out[blockIdx.x] = (int)s; + } +} +""" + + +def _drv(): + from cuda.bindings import driver + + return driver + + +def _ck(res, what: str): + """cuda-python returns (CUresult, *values): raise on an error, return the values.""" + err = res[0] if isinstance(res, tuple) else res + if int(err) != 0: + raise RuntimeError(f"locality: {what} failed: {err}") + if isinstance(res, tuple): + return res[1] if len(res) == 2 else res[1:] + return None + + +class _Mapping: + """One virtual range backed by per-domain physical allocations, exposed to torch through + ``__cuda_array_interface__``. Lives as long as a tensor that views it (the tensor holds this object).""" + + def __init__(self, device: int, pieces: Sequence[Tuple[int, int]]): + """``pieces``: (bytes, domain) in address order; each a multiple of the allocation granularity.""" + d = _drv() + self.device = device + self.size = sum(n for n, _ in pieces) + self.va = _ck(d.cuMemAddressReserve(self.size, 0, 0, 0), "cuMemAddressReserve") + self.handles = [] + off = 0 + try: + for nbytes, dom in pieces: + prop = _alloc_prop(device, dom) + h = _ck(d.cuMemCreate(nbytes, prop, 0), f"cuMemCreate(domain {dom})") + self.handles.append((h, nbytes)) + _ck(d.cuMemMap(int(self.va) + off, nbytes, 0, h, 0), "cuMemMap") + off += nbytes + acc = d.CUmemAccessDesc() + acc.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + acc.location.id = device + acc.flags = d.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _ck(d.cuMemSetAccess(self.va, self.size, [acc], 1), "cuMemSetAccess") + except Exception: + self.close() + raise + self.__cuda_array_interface__ = { + "shape": (self.size,), + "typestr": "|u1", + "data": (int(self.va), False), + "version": 3, + "strides": None, + "stream": None, + } + + def close(self): + d = _drv() + off = 0 + for h, n in self.handles: + d.cuMemUnmap(int(self.va) + off, n) + d.cuMemRelease(h) + off += n + self.handles = [] + if self.va is not None: + d.cuMemAddressFree(self.va, self.size) + self.va = None + + def __del__(self): + try: + self.close() + except Exception: + pass + + +_LIBCUDA = None + + +def _libcuda(): + global _LIBCUDA + if _LIBCUDA is None: + import ctypes + + _LIBCUDA = ctypes.CDLL("libcuda.so.1") + return _LIBCUDA + + +def _alloc_prop(device: int, domain: int): + d = _drv() + prop = d.CUmemAllocationProp() + prop.type = d.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.requestedHandleTypes = d.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE + prop.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN + prop.location.localized.deviceId = device + prop.location.localized.localityDomainId = domain + return prop + + +class LocalityMap: + """The locality domains of one device. Build with :meth:`get` before CUDA-graph capture.""" + + def __init__(self, device: int): + d = _drv() + _ck(d.cuInit(0), "cuInit") + dev = _ck(d.cuDeviceGet(device), "cuDeviceGet") + attr = d.CUdevice_attribute + self.device = device + self.n_domains = int( + _ck( + d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, dev), + "domain count", + ) + ) + self.sms_per_domain = int( + _ck( + d.cuDeviceGetAttribute( + attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_MULTIPROCESSOR_COUNT, dev + ), + "SMs per domain", + ) + ) + self.num_sms = int( + _ck(d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev), "SMs") + ) + if self.n_domains != 2: + raise RuntimeError( + f"locality: {self.n_domains} locality domains (this layer handles 2)" + ) + self._domain_sms = self._probe(d, dev) + dom = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for s in sms: + if dom[s] != -1: + raise RuntimeError(f"locality: SM {s} probed in two domains") + dom[s] = k + self.sm_domain: List[int] = dom + self.remainder: List[int] = [s for s in range(self.num_sms) if dom[s] < 0] + rt = _alloc_prop(device, 0) + self.granularity = int( + _ck( + d.cuMemGetAllocationGranularity( + rt, d.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM + ), + "granularity", + ) + ) + + @classmethod + def get(cls, device) -> "LocalityMap": + idx = torch.device(device).index if not isinstance(device, int) else device + idx = torch.cuda.current_device() if idx is None else idx + with _LOCK: + m = _MAPS.get(idx) + if m is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "LocalityMap probes the device: build it before CUDA-graph capture" + ) + m = _MAPS[idx] = cls(idx) + return m + + def sms(self, domain: int) -> List[int]: + return list(self._domain_sms[domain]) + + @property + def sm_table(self) -> torch.Tensor: + """int32 [num_sms] on the device: (domain << 16) | rank within the domain's sorted SMs; -1 = remainder SM. + For kernels that pick work by %smid themselves.""" + t = getattr(self, "_sm_table", None) + if t is None: + v = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for i, s in enumerate(sms): + v[s] = (k << 16) | i + t = self._sm_table = torch.tensor(v, dtype=torch.int32, device=f"cuda:{self.device}") + return t + + def _probe(self, d, dev) -> List[List[int]]: + """Split the SMs by locality domain (strict, no backfill) and record %smid of a probe kernel launched on a + green context of each part: the driver's own SM -> domain assignment.""" + from cuda.bindings import nvrtc + + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + # the probe module, compiled for this device + major = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev + ), + "cc", + ) + minor = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev + ), + "cc", + ) + prog = _ck( + nvrtc.nvrtcCreateProgram(_PROBE_SRC.encode(), b"smid_probe.cu", 0, [], []), + "nvrtcCreateProgram", + ) + opts = [f"--gpu-architecture=sm_{major}{minor}".encode()] + err = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)[0] + if int(err) != 0: + raise RuntimeError(f"locality: probe compile failed: {err}") + n = _ck(nvrtc.nvrtcGetCUBINSize(prog), "cubin size") + cubin = b" " * n + _ck(nvrtc.nvrtcGetCUBIN(prog, cubin), "cubin") + out: List[List[int]] = [] + nb = 4 * self.sms_per_domain # more blocks than SMs: every SM of the part shows up + torch.cuda.set_device(self.device) + buf = torch.full((nb,), -1, dtype=torch.int32, device=f"cuda:{self.device}") + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate(desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM), + "green ctx", + ) + st = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + ctx = _ck(d.cuCtxFromGreenCtx(g), "ctx from green ctx") + _ck(d.cuCtxPushCurrent(ctx), "push ctx") + try: + mod = _ck(d.cuModuleLoadData(cubin), "probe module") + fn = _ck(d.cuModuleGetFunction(mod, b"smid_probe"), "probe fn") + buf.fill_(-1) + torch.cuda.synchronize(self.device) + import ctypes + + arg = ctypes.c_void_p(buf.data_ptr()) + _ck( + d.cuLaunchKernel( + fn, nb, 1, 1, 32, 1, 1, 0, st, ((arg,), (ctypes.c_void_p,)), 0 + ), + "probe launch", + ) + _ck(d.cuStreamSynchronize(st), "probe sync") + sms = sorted(set(int(x) for x in buf.tolist() if x >= 0)) + d.cuModuleUnload(mod) + finally: + d.cuCtxPopCurrent() + d.cuStreamDestroy(st) + d.cuGreenCtxDestroy(g) + if len(sms) != self.sms_per_domain: + raise RuntimeError( + f"locality: domain {k} probed {len(sms)} SMs, expected {self.sms_per_domain}" + ) + out.append(sms) + return out + + def logical_map(self, per_domain: int, extra: int) -> Tuple[torch.Tensor, List[int]]: + """The SM -> logical block index permutation for a kernel of ``2 * per_domain + extra`` CTAs, one per SM: + logical blocks [d * per_domain, (d + 1) * per_domain) run on domain d's first ``per_domain`` SMs; logical + blocks 2 * per_domain .. run on the rest (the leftover domain SMs, alternating domains, then the remainder + SMs). Returns (int32 [num_sms] on the device, -1 = no block; the domain of each extra block, -1 remainder).""" + if 2 * per_domain + extra != self.num_sms: + raise ValueError( + f"locality: {2 * per_domain} + {extra} CTAs must cover the {self.num_sms} SMs" + ) + if per_domain > self.sms_per_domain: + raise ValueError(f"locality: {per_domain} CTAs per domain > {self.sms_per_domain} SMs") + lmap = [-1] * self.num_sms + for k in range(2): + for i, s in enumerate(self._domain_sms[k][:per_domain]): + lmap[s] = k * per_domain + i + left = [self._domain_sms[k][per_domain:] for k in range(2)] + rest = [x for pair in zip(*left) for x in pair] + self.remainder + ext_dom = [] + for i, s in enumerate(rest): + lmap[s] = 2 * per_domain + i + ext_dom.append(self.sm_domain[s]) + assert sorted(lmap) == list(range(self.num_sms)) + return torch.tensor(lmap, dtype=torch.int32, device=f"cuda:{self.device}"), ext_dom + + def domain_streams(self) -> List[torch.cuda.ExternalStream]: + """One non-blocking stream per domain on a green context holding exactly that domain's SMs (kept alive with + this map): work launched on stream d runs only on domain d's SMs. Fork / join with events from the caller's + stream; graph capture follows the forks. Built once, before capture.""" + st = getattr(self, "_streams", None) + if st is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.domain_streams: build before CUDA-graph capture") + d = _drv() + dev = _ck(d.cuDeviceGet(self.device), "cuDeviceGet") + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + self._green = [] + st = [] + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate( + desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM + ), + "green", + ) + s_ = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + self._green.append((g, s_)) + st.append(torch.cuda.ExternalStream(int(s_), device=f"cuda:{self.device}")) + self._streams = st + return st + + def domain_of(self, ptr: int) -> int: + """The locality domain ordinal of a device address (-1: not placed in one domain). Through libcuda + directly: cuda-python's attribute helper does not know this attribute.""" + import ctypes + + d = _drv() + lib = _libcuda() + out = ctypes.c_int(-9) + attr = int(d.CUpointer_attribute.CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL) + err = lib.cuPointerGetAttribute(ctypes.byref(out), ctypes.c_int(attr), ctypes.c_uint64(ptr)) + if err != 0: + raise RuntimeError( + f"locality: cuPointerGetAttribute(LOCALITY_DOMAIN_ORDINAL) failed: {err}" + ) + return int(out.value) + + def place_split(self, t: torch.Tensor, split_row: int) -> torch.Tensor: + """Split-mapped placement: a copy of ``t`` (contiguous, rows along dim 0) with the SAME shape and layout, in one + virtual range whose rows [0, split_row) are physically in domain 0 and rows [split_row, N) in domain 1. The + split must fall on an allocation granule (rows [0, split_row) a multiple of 2 MiB on VR200); the tail is padded + to a granule in domain 1. Stock consumers read it as a plain tensor; no extra memory once ``t`` is dropped.""" + if not t.is_contiguous(): + raise ValueError("locality.place_split: needs a contiguous tensor") + row_bytes = t[0].numel() * t.element_size() if t.dim() > 1 else t.element_size() + head = split_row * row_bytes + total = t.numel() * t.element_size() + if not 0 < head < total or head % self.granularity: + raise ValueError( + f"locality.place_split: rows [0, {split_row}) = {head} B must be a granule multiple " + f"({self.granularity} B) inside the {total} B tensor" + ) + tail = -(-(total - head) // self.granularity) * self.granularity + m = _Mapping(self.device, [(head, 0), (tail, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + out = flat[:total].view(t.dtype).view(t.shape) + out.copy_(t) + out._locality_mapping = m + for q, k in ( + (out.data_ptr(), 0), + (out.data_ptr() + head - 1, 0), + (out.data_ptr() + head, 1), + (out.data_ptr() + total - 1, 1), + ): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.place_split: byte {q:#x} is not in domain {k}") + return out + + def replace_param_(self, param: torch.nn.Parameter, split_row: int) -> torch.nn.Parameter: + """Re-place a parameter in place with :meth:`place_split` (swaps ``param.data``; the old block returns to the + caching allocator, so call ``torch.cuda.empty_cache()`` between large re-placements).""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.replace_param_: not during CUDA-graph capture") + param.data = self.place_split(param.data, split_row) + return param + + def split_rows(self, t: torch.Tensor, dim: int = 1) -> torch.Tensor: + """Place ``t`` (contiguous, 1-byte elements, rows along ``dim`` split into two equal halves) as + [2, *t.shape[:dim], t.shape[dim] // 2, *t.shape[dim + 1:]]: half k (rows [k h, (k + 1) h) of every leading + index) in domain k's memory, one virtual range. Each half's byte size must be a multiple of the allocation + granularity (2 MiB on VR200). The copy is exact; the caller drops ``t`` afterwards.""" + if t.element_size() != 1 or not t.is_contiguous(): + raise ValueError("locality.split_rows: needs a contiguous 1-byte tensor") + n = t.shape[dim] + if n % 2: + raise ValueError(f"locality.split_rows: {n} rows don't split in halves") + half_bytes = t.numel() // 2 + if half_bytes % self.granularity: + raise ValueError( + f"locality.split_rows: a half is {half_bytes} B, not a multiple of the {self.granularity} B granularity" + ) + m = _Mapping(self.device, [(half_bytes, 0), (half_bytes, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + shape = (2,) + tuple(t.shape[:dim]) + (n // 2,) + tuple(t.shape[dim + 1 :]) + out = flat.view(torch.uint8).view(shape) + src = t.view(torch.uint8) + for k in range(2): + out[k].copy_(src.narrow(dim, k * (n // 2), n // 2)) + out._locality_mapping = m # the tensor keeps its mapping alive + for k in range(2): + p = out[k].data_ptr() + for q in (p, p + half_bytes - 1): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.split_rows: byte {q:#x} is not in domain {k}") + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_direct.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_direct.py new file mode 100644 index 000000000000..30b905712b1f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_direct.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +"""Register-prefetch BF16 GEMM for low-M, long-K decode shapes. + +This target's copy of the installed FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_direct.py`` with one +addition, the drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH_DIRECT): after its one store of each output, +the owning thread also stores the CTA's column pairs as the boundary all-reduce's canonical exchange words (kernels/ +m3_mc_allreduce) at the call parity, before the kernel's (unchanged, last) release of its dependents. The plain +kernel and its output are FlashInfer's; only the push variant is called here (``run_direct_dense_push``). + +The kernel keeps a complete output dot product inside one CTA and reuses each +prefetched B value across several public-M rows. It is intentionally a +separate autotuner runner from the Blackwell tensor-core split-K kernel: the +two algorithms have different useful shape regions and tactic spaces. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import torch as _torch +from cutlass import Int32, const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +_VECTOR_WIDTH = 8 # capacity: fixed: a vector load width +_SUPPORTED_BLOCK_SIZES = (32, 64, 96, 128, 192, 256, 384) +_SUPPORTED_OUTPUTS_PER_BLOCK = (1, 2, 4) +_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.s32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + llvm.inline_asm( + None, + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dataclasses.dataclass(frozen=True) +class DirectTactic: + """One direct-kernel specialization.""" + + block_size: int + outputs_per_block: int + rows_per_block: int + + +def _default_rows_per_block(m: int) -> int: + if m <= 8: + return m + return next(rows for rows in (8, 4, 2, 1) if m % rows == 0) + + +def validate_tactic(tactic: DirectTactic, m: int, n: int, k: int) -> None: + """Reject a direct tactic that cannot serve ``(m, n, k)``.""" + if tactic.block_size not in _SUPPORTED_BLOCK_SIZES: + raise ValueError(f"unsupported block_size={tactic.block_size}") + if tactic.outputs_per_block not in _SUPPORTED_OUTPUTS_PER_BLOCK: + raise ValueError(f"unsupported outputs_per_block={tactic.outputs_per_block}") + if not 1 <= m <= _MAX_M: + raise ValueError(f"direct GEMM requires 1 <= M <= {_MAX_M}, got {m}") + if not 1 <= tactic.rows_per_block <= m or m % tactic.rows_per_block: + raise ValueError(f"rows_per_block={tactic.rows_per_block} must divide M={m}") + if n <= 0 or n % tactic.outputs_per_block: + raise ValueError(f"N={n} must be divisible by outputs_per_block={tactic.outputs_per_block}") + k_tile = tactic.block_size * _VECTOR_WIDTH + if k <= 0 or k % k_tile: + raise ValueError(f"K={k} must be divisible by {k_tile}") + + +def default_tactic(m: int, n: int, k: int) -> DirectTactic: + """Choose the measured register-prefetch fallback tactic.""" + block_size = next( + (block for block in (256, 192, 128, 96, 64, 32) if k % (block * _VECTOR_WIDTH) == 0), + None, + ) + if block_size is None: + raise ValueError("direct GEMM requires a supported 16-byte K tiling") + outputs_per_block = next(outputs for outputs in (2, 1) if n % outputs == 0) + tactic = DirectTactic( + block_size, + outputs_per_block, + _default_rows_per_block(m), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +def autotune_tactics(m: int, n: int, k: int) -> list[DirectTactic]: + """Enumerate the compact tactic space used by FlashInfer autotuning. + + Block sizes cover every configuration exercised in the H100/B200 sweep; + output grouping spans the measured 1/2/4-column choices. Row tiling stays + at the occupancy-oriented default to keep JIT cost bounded. + """ + try: + default = default_tactic(m, n, k) + except ValueError: + return [] + tactics = [default] + for block_size in _SUPPORTED_BLOCK_SIZES: + for outputs_per_block in _SUPPORTED_OUTPUTS_PER_BLOCK: + tactic = DirectTactic( + block_size, + outputs_per_block, + default.rows_per_block, + ) + try: + validate_tactic(tactic, m, n, k) + except ValueError: + continue + tactics.append(tactic) + return list(dict.fromkeys(tactics)) + + +def prefer_direct_bf16_gemm_sm100(m: int, n: int, k: int) -> bool: + """Return the conservative B200 no-autotune crossover heuristic. + + The three bands are a compact fit to a warm/cold sweep over M=1..16,24,32, + 18 N values, and 11 K values. This is deliberately not a blanket rule for + K=8192: direct wins only where public M and N leave the tensor-core path + with too little independent output work. + """ + return k == 8192 and ((m == 1 and n <= 4608) or (m <= 4 and n <= 512) or (m <= 8 and n <= 256)) + + +class DirectDenseGemmKernel: + """K-specialized direct GEMM with whole-mainloop vector prefetch.""" + + def __init__( + self, + *, + element_type, + num_rows: int, + k_extent: int, + tactic: DirectTactic, + use_pdl: bool, + ar_push: bool = False, + ) -> None: + validate_tactic(tactic, num_rows, tactic.outputs_per_block, k_extent) + self.ar_push = ar_push + self.element_type = element_type + self.num_rows = num_rows + self.rows_per_block = tactic.rows_per_block + self.k_extent = k_extent + self.block_size = tactic.block_size + self.outputs_per_block = tactic.outputs_per_block + self.vector_width = _VECTOR_WIDTH + self.use_pdl = use_pdl + self.num_warps = tactic.block_size // cute.arch.WARP_SIZE + self.num_k_tiles = k_extent // (tactic.block_size * _VECTOR_WIDTH) + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64 = cutlass.Int64(0), + ar_flags: cutlass.Int64 = cutlass.Int64(0), + ar_rank: Int32 = Int32(0), + ar_world: Int32 = Int32(1), + ar_maxt: Int32 = Int32(8), + ) -> None: + n = cute.size(gB, mode=[0]) + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel(gA, gB, gC, copy_a, copy_b, ar_ws, ar_flags, ar_rank, ar_world, ar_maxt).launch( + grid=[ + cute.ceil_div(n, self.outputs_per_block), + self.num_rows // self.rows_per_block, + 1, + ], + block=[self.block_size, 1, 1], + smem=self.rows_per_block * self.outputs_per_block * self.num_warps * 4, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + copy_a: cute.CopyAtom, + copy_b: cute.CopyAtom, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, block_m, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + + num_rows: cutlass.Constexpr = self.rows_per_block + outputs_per_block: cutlass.Constexpr = self.outputs_per_block + vector_width: cutlass.Constexpr = self.vector_width + block_size: cutlass.Constexpr = self.block_size + num_warps: cutlass.Constexpr = self.num_warps + num_k_tiles: cutlass.Constexpr = self.num_k_tiles + + acc = cute.make_rmem_tensor( + cute.make_layout((num_rows, outputs_per_block), stride=(outputs_per_block, 1)), + cutlass.Float32, + ) + acc.fill(0.0) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + n_base = block_idx * outputs_per_block + m_base = block_m * num_rows + gA_vec = cute.logical_divide(gA, (None, vector_width)) + gB_vec = cute.logical_divide(gB, (None, vector_width)) + tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) + tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) + tA = tA_all[None, (None, (tidx, None))] + + b_regs = cute.make_rmem_tensor( + cute.make_layout( + (outputs_per_block, num_k_tiles, vector_width), + stride=(num_k_tiles * vector_width, vector_width, 1), + ), + self.element_type, + ) + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + for k_tile in cutlass.range_constexpr(num_k_tiles): + cute.copy(copy_b, tB[None, k_tile], b_regs[ni, k_tile, None]) + + a_regs = cute.make_rmem_tensor( + cute.make_layout((num_k_tiles, vector_width), stride=(vector_width, 1)), + self.element_type, + ) + for mi in cutlass.range_constexpr(num_rows): + for k_tile in cutlass.range_constexpr(num_k_tiles): + cute.copy( + copy_a, + tA[m_base + mi, None, k_tile], + a_regs[k_tile, None], + ) + for k_tile in cutlass.range_constexpr(num_k_tiles): + for vi in cutlass.range_constexpr(vector_width): + a_value = a_regs[k_tile, vi].to(cutlass.Float32) + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_value * b_regs[ni, k_tile, vi].to( + cutlass.Float32 + ) + + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + smem_layout = cute.make_layout( + (num_rows, outputs_per_block, num_warps), + stride=(outputs_per_block * num_warps, num_warps, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + partials[mi, ni, warp_idx] = acc[mi, ni] + + cute.arch.sync_threads() + if tidx == 0: + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + total = cutlass.Float32(0.0) + for warp in cutlass.range_constexpr(num_warps): + total = total + partials[mi, ni, warp] + gC[m_base + mi, n_base + ni] = total.to(self.element_type) + if const_expr(self.ar_push): + acc[mi, ni] = total + if const_expr(self.ar_push): + # Push the CTA's column pairs of every row into the boundary all-reduce's exchange (prepushed 2), at the + # (token, column / (N / 8)) call parity; the release below follows the pushes (the prepushed contract). + n_cols = cute.size(gC, mode=[1]) + half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(n_cols // 2) + for mi in cutlass.range_constexpr(num_rows): + t = m_base + mi + for pi in cutlass.range_constexpr(outputs_per_block // 2): + col = n_base + 2 * pi + slot = cutlass.Int64(t * 8 + col // (n_cols // 8)) + par = _ld_volatile_s32(ar_flags + cutlass.Int64(4) * slot) & Int32(1) + row = cutlass.Int64(t * ar_world + ar_rank) + off = ( + cutlass.Int64(par) * half + + row * cutlass.Int64(n_cols // 2) + + cutlass.Int64(col // 2) + ) + word = _bf16x2_word(acc[mi, 2 * pi], acc[mi, 2 * pi + 1]) + _st_global_b32(ar_ws + off * cutlass.Int64(4), word) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +def _from_dlpack_static(tensor: _torch.Tensor): + # K is specialized and the row stride must retain its 16-byte divisibility + # for the verifier to accept vectorized G2R copies. + return from_dlpack(tensor, assumed_align=32) + + +def _make_compile_repr_tensors(dtype, m: int, n: int, k: int): + return tuple( + _from_dlpack_static(tensor) + for tensor in ( + _torch.empty((m, k), dtype=dtype, device="cuda"), + _torch.empty((n, k), dtype=dtype, device="cuda"), + _torch.empty((m, n), dtype=dtype, device="cuda"), + ) + ) + + +@functools.cache +def _get_compiled_direct_kernel( + dtype, + m: int, + n: int, + k: int, + tactic: DirectTactic, + use_pdl: bool, +): + if dtype != _torch.bfloat16: + raise ValueError(f"direct GEMM supports BF16; got {dtype}") + kernel = DirectDenseGemmKernel( + element_type=cutlass.BFloat16, + num_rows=m, + k_extent=k, + tactic=tactic, + use_pdl=use_pdl, + ) + tensors = _make_compile_repr_tensors(dtype, m, n, k) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + return cute_ext.compile(kernel, *tensors, stream, options=_COMPILE_OPTIONS) + + +def _validate_runtime_tensors(a, b, out, tactic: DirectTactic): + if any(not isinstance(tensor, _torch.Tensor) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("direct GEMM accepts only 2D tensors") + if a.device.type != "cuda" or b.device != a.device or out.device != a.device: + raise ValueError("a, b, and out must be on the same CUDA device") + if a.dtype != _torch.bfloat16 or b.dtype != a.dtype or out.dtype != a.dtype: + raise ValueError("a, b, and out must share BF16 dtype") + if not a.is_contiguous() or not b.T.is_contiguous() or not out.is_contiguous(): + raise ValueError("direct GEMM requires row-major A/out and column-major B") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + validate_tactic(tactic, m, n, k) + return m, n, k + + +def run_direct_dense(a, b, out, pdl: bool, tactic: DirectTactic): + """Run direct ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + m, n, k = _validate_runtime_tensors(a, b, out, tactic) + compiled = _get_compiled_direct_kernel(a.dtype, m, n, k, tactic, pdl) + tensors = tuple(_from_dlpack_static(tensor) for tensor in (a, b.T, out)) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled(*tensors, stream) + return out + + +__all__ = [ + "DirectTactic", + "autotune_tactics", + "default_tactic", + "prefer_direct_bf16_gemm_sm100", + "run_direct_dense", + "run_direct_dense_push", + "compile_direct_push", +] + + +@functools.cache +def _get_compiled_direct_push(dtype, m: int, n: int, k: int, tactic: DirectTactic, use_pdl: bool): + assert tactic.outputs_per_block % 2 == 0, "the push writes column pairs" + kernel = DirectDenseGemmKernel( + element_type=cutlass.BFloat16, + num_rows=m, + k_extent=k, + tactic=tactic, + use_pdl=use_pdl, + ar_push=True, + ) + tensors = _make_compile_repr_tensors(dtype, m, n, k) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kernel, + *tensors, + stream, + cutlass.Int64(0), + cutlass.Int64(0), + Int32(0), + Int32(1), + Int32(8), + options=_COMPILE_OPTIONS, + ) + + +def compile_direct_push(a, b, out, pdl: bool, tactic) -> None: + """Compile (only) the push variant for these operands' shapes.""" + m, n, k = _validate_runtime_tensors(a, b, out, DirectTactic(*tactic)) + _get_compiled_direct_push(a.dtype, m, n, k, DirectTactic(*tactic), pdl) + + +def run_direct_dense_push(a, b, out, pdl: bool, tactic: DirectTactic, ar: tuple): + """``run_direct_dense`` that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange.""" + m, n, k = _validate_runtime_tensors(a, b, out, tactic) + compiled = _get_compiled_direct_push(a.dtype, m, n, k, tactic, pdl) + tensors = tuple(_from_dlpack_static(tensor) for tensor in (a, b.T, out)) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + ws, flags, rank, world, maxt = ar + compiled( + *tensors, + stream, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_splitk.py new file mode 100644 index 000000000000..466a327a4bbf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_splitk.py @@ -0,0 +1,1249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +This target's copy of the installed FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the +Eagle3 drafter's low-M linears, modeling's ``DRAFT_LOWM_OWN_SPLITK``): the same kernel and the same split order at every +tactic, bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's +final cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, +the SwiGLU prologue that follow. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 # capacity: fixed: a kernel tile constant + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 # capacity: fixed: an alignment + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 # capacity: fixed: pipeline stages +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 # capacity: fixed: pipeline stages + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= _MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES + base = SplitKTactic(mma_m, mma_n, split_k, min_stages) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(min_stages, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 # capacity: fixed: a split index, not a size + + +# The drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH): the exchange words of +# kernels/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.s32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + llvm.inline_asm( + None, + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + ar_push: bool = False, + ) -> None: + self.ar_push = ar_push + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64 = cutlass.Int64(0), + ar_flags: cutlass.Int64 = cutlass.Int64(0), + ar_rank: Int32 = Int32(0), + ar_world: Int32 = Int32(1), + ar_maxt: Int32 = Int32(8), + ar_tokens: Int32 = Int32(0), + ar_n: Int32 = Int32(0), + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel( + a, b, c, bias, ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_tokens, ar_n + ).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + ar_ws: cutlass.Int64, # ar_push: the boundary all-reduce's multicast exchange words (int32) + ar_flags: cutlass.Int64, # ar_push: its per-(token, cluster CTA) call flags (parity = flag & 1) + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, # the bucket's tokens + ar_tokens: Int32, # rows pushed (public M) + ar_n: Int32, # public N (the hidden size): the exchange row width + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + # ar_push: the owner's BF16 tile (kernel-M = output column, kernel-N = token), M-major. + sD = cute_ext.allocate( + mC.element_type, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=16, + ) + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sD, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl and not self.ar_push): + # ar_push: the dependent is a prepushed all-reduce, which skips its grid wait: it is released only after + # this grid's own grid wait and its pushes, in the owner's _store_tile (the prepushed >= 2 contract). + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sD: cute.Tensor, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + self._store_tile(thr_t2r, rAcc, rD, gD_epi, c_dtype) + if cutlass.const_expr(self.ar_push): + self._push_tile( + thr_t2r, + rD, + sD, + epi_tid, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + @cute.experimental.jit + def _store_tile(self, thr_t2r, rAcc, rD, gD_epi, c_dtype: cutlass.Constexpr): + """The owner rank's reduced FP32 tile: one cast to ``c_dtype`` and the store to global memory.""" + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + @cute.experimental.jit + def _push_tile( + self, + thr_t2r, + rD, + sD, + epi_tid: cutlass.Int32, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output + would store them (kernels/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" + cute_ext.partition_and_copy(thr_t2r, rD, sD) + cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + pairs = self.cta_m // 2 + col0 = bidx * self.cta_m + half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) + for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): + word = epi_tid + w0 + t = word // pairs + p = word % pairs + if t < ar_tokens: + n = col0 + 2 * p + par = _ld_volatile_s32( + ar_flags + cutlass.Int64(4) * cutlass.Int64(t * 8 + n // (ar_n // 8)) + ) & Int32(1) + off = ( + cutlass.Int64(par) * half + + cutlass.Int64(t * ar_world + ar_rank) * cutlass.Int64(ar_n // 2) + + cutlass.Int64(n // 2) + ) + _st_global_b32( + ar_ws + off * cutlass.Int64(4), + _bf16x2_word( + sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + ), + ) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +import torch as _torch # noqa: E402 (as in FlashInfer) + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ) + + +@cute.experimental.jit +def _bmm_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ar_tokens: Int32, + ar_n: Int32, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any(tensor.dtype != a.dtype for tensor in tensors): + raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a, + b, + bias, + out, + pdl: bool, + tactic: SplitKTactic, +): + """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +@functools.cache +def _get_compiled_push( + dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] +): + kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, ar_push=True) + compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + z64, z32 = cutlass.Int64(0), Int32(0) + return cute_ext.compile( + _bmm_push, + kernel, + *compile_tensors[:3], + stream, + z64, + z64, + z32, + Int32(1), + Int32(8), + z32, + Int32(6144), + ) + + +def compile_push(a, b, out, pdl: bool, tactic) -> None: + """Compile (only) the push variant of ``run_splitk_dense_push`` for these operands' shapes and layouts.""" + cute_tensors = _to_cute_swap(a, b, out, None) + _get_compiled_push(a.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) + + +def run_splitk_dense_push(a, b, out, pdl: bool, tactic: SplitKTactic, ar: tuple): + """``run_splitk_dense`` without bias that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange; + ``ar`` = (multicast words address, flags address, rank, world, bucket tokens) of the 8-token bucket.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_push(a.dtype, tactic, pdl, cute_tensors[4]) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + ws, flags, rank, world, maxt = ar + compiled( + *cute_tensors[:3], + stream, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + Int32(out.shape[0]), + Int32(out.shape[1]), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_mc_allreduce.py new file mode 100644 index 000000000000..76b73fc9ce91 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_mc_allreduce.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-latency TP all-reduce + residual + RMSNorm for decode (M <= 8 tokens) over a multicast mapping. + + partial = bf16(x + x2) (optional second input, e.g. shared + routed experts) + updated = bf16(residual + sum_r partial_r) (ranks summed in FP32, rank order) + normed = bf16(updated * rsqrt(mean(updated^2) + eps) * w) + +One cluster of CLUSTER CTAs per token; CTA c owns columns [c * H / CLUSTER, ...), one 16-byte chunk (8 BF16) +per thread. Each CTA pushes its slice of this rank's partial through the multicast mapping (one store lands in +every rank's buffer), releases dependents, polls its own buffer until all ranks' slices arrived, sums them, adds +the residual and reduces the RMS statistic across the cluster through DSMEM. + +Exchange protocol as in the K3 sandwich: int32 words, 0x80000000 = empty (pushed -0.0 halves become +0.0), each +CTA keeps a call counter whose parity selects the buffer half, readers re-arm the words they read. All ranks of +the TP group must issue the same calls in the same stream order. + +The workspace holds two exchange sets. With prepushed 3 two producers pushed this rank's x (set A) and x2 (set B), +e.g. the routed experts' finalize and the shared expert's down projection, and the kernel forms bf16(x + x2) per +rank from the two slots, exactly as the unpushed path does. + +Producers that push (the o_proj / shared-expert GEMV ``ar`` tuple, the routed-expert engine's ``arp`` descriptor) are +pointed at a bucket's exchange by its ``mc`` address, ``flags_for(8)`` counters, ``rank``, ``world_size`` and +``max_tokens``; ``get_workspace`` and ``_bucket`` select it. + +The kernel is ``torch.ops.trtllm.minimax_m3_mc_allreduce_norm``, reached through the ``ops/comm`` entries +``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``; this module owns the workspaces, the bucket and +epilogue-option choice and the MXFP8 publication. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op +from ..ops.comm.m3_mc_finalize_allreduce_norm import ( + m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, +) + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +BUCKETS = (8, 16, 48) # token capacity of each workspace (grid = bucket * cluster CTAs) +# The MoE finalize (cluster kernel) stops at FINALIZE_MAX_TOKENS (= MAX_TOKENS here). +FINALIZE_MAX_TOKENS = 48 # capacity: fixed: the finalize kernel's built rows +EMPTY_WORD = -2147483648 + +# Attention / MoE boundary epilogue options (each bit-identical): 1 residual + norm weight loaded at entry (before +# the poll); 2 cluster RMS exchange by st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated +# right after the poll (no trailing __syncthreads); 8 max-smem carveout so the next kernel's CTAs co-reside early. +# The 8-token bucket, by token count 1..8 (GB300 TP4 tuning). +AR_OPTS_BY_M = (7, 7, 15, 15, 15, 15, 7, 7) +# The 16- and 48-token buckets. +AR_OPTS_BY_BUCKET = {16: 15, 48: 15} +# The MoE-finalize all-reduce, per bucket. +AR_FIN_OPTS_BY_BUCKET = {8: 7, 16: 15, 48: 15} + + +# Option 2 is also the safe ordering of the cluster RMS exchange: its stores into peer CTAs' shared memory follow a +# cluster barrier arrive / wait pair. Without it they run before any cluster barrier, when a peer need not have +# started, so every call keeps it. +_OPT_CLUSTER_BARRIER = 2 + + +def _ar_opts(m: int, bucket: int, finalize: bool = False) -> int: + """Epilogue options for an m-row call on the given bucket; every option is bit-identical. A bucket missing from + the tables gets 0, and every result keeps _OPT_CLUSTER_BARRIER.""" + if finalize: + opts = AR_FIN_OPTS_BY_BUCKET.get(bucket, 0) + elif bucket == 8: + opts = AR_OPTS_BY_M[m - 1] + else: + opts = AR_OPTS_BY_BUCKET.get(bucket, 0) + return opts | _OPT_CLUSTER_BARRIER + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags + + +class McBuckets: + """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" + + def __init__(self, workspaces): + self.workspaces = workspaces + self.hidden = workspaces[0].hidden + + def for_tokens(self, m: int): + for ws in self.workspaces: + if m <= ws.max_tokens: + return ws + return None + + +def _bucket(workspace, m): + return workspace.for_tokens(m) if isinstance(workspace, McBuckets) else workspace + + +def decline_reason(x, residual, workspace, cluster=8, x2=None) -> str: + """Why ``m3_mc_allreduce_norm`` returned None for these operands (for the dispatch's miss message).""" + if workspace is None: + return "no multicast workspace: it is created collectively on the first eager decode step" + bucket = _bucket(workspace, x.shape[0]) + if bucket is None: + return f"no workspace bucket holds {x.shape[0]} tokens (buckets {BUCKETS})" + if bucket.flags.get(cluster) is None: + return f"no {cluster}-CTA call counters for the {bucket.max_tokens}-token bucket" + return ( + f"operands outside the kernel: x {tuple(x.shape)} {x.dtype} contiguous={x.is_contiguous()}, " + f"residual contiguous={residual.is_contiguous()}, hidden {bucket.hidden}" + + ("" if x2 is None else f", x2 {tuple(x2.shape)} contiguous={x2.is_contiguous()}") + ) + + +def m3_mc_allreduce_norm( + x, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + x2=None, + quant=False, + prepushed=0, + drafter_norm=False, + capture=None, +): + """Return (normed, updated_residual) or None when unsupported. ``drafter_norm``: the Eagle3 drafter boundary + (opts bit 16: normalize the unrounded FP32 sum, as the drafter's stock FlashInfer fused_add_rmsnorm does); 8-token + bucket only, x only, no MXFP8 copy. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture`` = (out [>= M, H] view, fc_norm weight + [H], eps): the boundary also writes the Eagle3 capture of the updated residual, bit for bit + m3_spec_tail.capture_norm (CAPTURE_NORM_IN_BOUNDARY; not with ``drafter_norm``). + + A None makes the caller take another collective, so it must be the same on every TP rank. It is: it depends only + on the shapes, dtypes and layouts of the operands (the same code builds them on every rank) and on whether the + workspace's flags exist, which ``get_workspace`` creates collectively, before any capture, for cluster 8.""" + workspace = _bucket(workspace, x.shape[0]) + if workspace is None: + return None + if ( + x.dim() != 2 + or x.shape[0] > MAX_TOKENS + or x.dtype != torch.bfloat16 + or x.shape[1] != workspace.hidden + or not x.is_contiguous() + or not residual.is_contiguous() + or (x2 is not None and (not x2.is_contiguous() or x2.shape != x.shape)) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + if drafter_norm and ( + workspace.max_tokens != 8 or x2 is not None or quant or prepushed not in (0, 2) + ): + return None + outs = _allreduce_norm_op( + x, + x2, + residual, + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and x.shape[1] % 128 == 0), + int(prepushed), + _ar_opts(x.shape[0], workspace.max_tokens) | (16 if drafter_norm else 0), + *_capture_args(capture, drafter_norm), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def m3_mc_finalize_allreduce_norm( + fc2, + perm, + scales, + shared, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + quant=False, + capture=None, +): + """MoE finalize (sum_k scale * fc2[perm] + shared) + all-reduce + residual + RMSNorm in one kernel. + Returns (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture``: as ``m3_mc_allreduce_norm``.""" + m = residual.shape[0] + workspace = _bucket(workspace, m) + if workspace is None: + return None + if ( + residual.dim() != 2 + or m > MAX_TOKENS + or residual.shape[1] != workspace.hidden + or fc2.dtype != torch.bfloat16 + or not fc2.is_contiguous() + or fc2.dim() != 2 + or perm.dtype != torch.int32 + or perm.dim() != 2 + or perm.shape[0] != m + or not perm.is_contiguous() + or scales.dtype not in (torch.float32, torch.bfloat16) + or not scales.is_contiguous() + or (shared is not None and (shared.shape != residual.shape or not shared.is_contiguous())) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _finalize_allreduce_norm_op( + fc2, + perm, + scales, + shared, + residual.contiguous(), + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and residual.shape[1] % 128 == 0), + _ar_opts(m, workspace.max_tokens, finalize=True), + *_capture_args(capture, False), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def _capture_args(capture, drafter_norm: bool) -> tuple: + """(cap_out, cap_w, cap_eps) for the op, from ``capture`` = (out, fc_norm weight, eps) or None.""" + if capture is None: + return (None, None, 0.0) + assert not drafter_norm, "the capture epilogue is not built for the drafter boundary" + out, weight, eps = capture + return (out, weight, float(eps)) + + +_WORKSPACE: dict = {} +_AGREE_COMMS: dict = {} + + +def _tp_comm(mapping): + """A communicator over mapping's TP group, created once (collective over that group only).""" + key = tuple(mapping.tp_group) + comm = _AGREE_COMMS.get(key) + if comm is None: + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + comm = mapping.tp_group_pg + else: + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + _AGREE_COMMS[key] = comm + return comm + + +def tp_agree(mapping, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of mapping's TP group (one blocking host collective). + + For a fast path whose availability is decided from rank-local state (a kernel build that can fail on one rank): + every rank must take the same path, so that they issue the same collectives with the same exchange protocol and + the same kernels. Call it once per decision, at build time, at the same point on every rank, outside CUDA-graph + capture; never per step. With one rank, or attention DP (whose ranks run different batches), it returns ``ok``. + """ + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return ok + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("tp_agree must run outside CUDA-graph capture") + from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded + + return _mnnvl_workspace_all_succeeded(_tp_comm(mapping), bool(ok)) + + +def tp_agree_built(mapping, built, what: str): + """``built`` when it is not None on every TP rank, else None on every rank (see ``tp_agree``).""" + if tp_agree(mapping, built is not None) or built is None: + return built + from tensorrt_llm.logger import logger + + logger.warning(f"{what}: built here but not on another TP rank; every rank falls back") + return None + + +def get_workspace(mapping, hidden: int): + """Process-wide workspace (created collectively on first eager use); None while capturing before creation. + + The first eager call for a hidden size (the warmup forward, before any CUDA-graph capture) creates one + McAllReduceWorkspace per bucket in BUCKETS on every TP rank (MPI collective over the TP group) plus each + bucket's 8-CTA-cluster call counters; later calls, captured or not, return the same McBuckets. The workspaces + live until process exit. + """ + ws = _WORKSPACE.get(hidden) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + return None + wss = [McAllReduceWorkspace.create(mapping, hidden, b) for b in BUCKETS] + for w in wss: + w.flags_for(8) + ws = McBuckets(wss) + _WORKSPACE[hidden] = ws + return ws diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/__init__.py new file mode 100644 index 000000000000..2ed197d5c475 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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 sm_107 TP2 target for low-batch decode as a CuTe DSL kernel: ``m3_moe`` for 1..8 tokens +(op.py) at the TP2 routed and shared intermediate of 1536 per rank.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..084f54951209 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3936 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP2 rank (routed and shared intermediate 1536): +``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 # capacity: fixed: the checkpoint's hidden size +# routed intermediate per rank: 1536 at TP2 (this target: FC2 K = 6 full 256-K tiles), 768 at TP4, 384 at TP8 (FC2 K +# = 1.5 256-K tiles: the last tile's A box is half out of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 +# MMAs / scale atoms are issued; no padding, no extra bytes) +I_TP = _cfg("i_tp", 768) +TWO_I = 2 * I_TP +assert I_TP % 128 == 0 and 128 <= I_TP <= 1536, I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 # capacity: fixed: the CTA size +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# TAIL_TRIGGER (with LATE_TRIGGER, no all-reduce push): the CTAs that own more FC2 blocks than the rest (the first +# H / 32 mod NCTA, which end last) release the dependents right after the grid dependency, so the dependents launch +# when the other CTAs tear down, onto the SMs they leave, instead of after the last of the heavy CTAs. +TAIL_TRIGGER = bool(_cfg("tail_trigger", 1)) +NHEAVY = (H // 32) % NCTA if NCTA < H // 32 else 0 +# SH_TAKE: the last NHEAVY shared-expert CTAs, after their shared work, each run the routed body as an FC2-only CTA for +# one FC2 output block (the second block of routed CTA c < NHEAVY, which then owns one). The plan carries their rows +# (PLAN_COLS columns). Each block keeps its tiles, scales and finalize, so the output is unchanged. +SH_TAKE = bool(_cfg("sh_take", 0)) and SHARED and 0 < NHEAVY <= 24 +SH_HELPER0 = 24 - NHEAVY # first helper shared index (NSH = 24) +PLAN_COLS = NCTA + 24 if SH_TAKE else NCTA +# EARLY_AR (scheduling only): when the finalize pushes into the multicast all-reduce, EVERY CTA releases the +# dependents right after the grid dependency, so the all-reduce (which reads only the pushed tagged words, without a +# grid wait) and, through its own early release, the next layer's QKV GEMV become resident on each SM this kernel +# frees, instead of after the last CTA's teardown. The arithmetic is unchanged. +EARLY_AR = bool(_cfg("early_ar", 1)) +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +# the shared intermediate per rank: 1536 (TP2), 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate +# positions: its gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP2: 1 x (64 + +# 64), full K in one accumulator, the order of the stock tcgen05 MXFP8 GEMM; TP4: 2 x (32 + 32), the order of +# m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. +SH_I = _cfg("sh_i", 768) +SH_ROWS = SH_I // 24 +# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) +SH_SPLIT = 128 // (2 * SH_ROWS) +SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N +SH_KB = H // 32 # capacity: fixed: a kernel tile constant (gate_up K blocks) +SH_DKB = SH_I // 32 # down K blocks +SH_DKT = SH_I // 128 # down k-tiles +SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the +# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction +# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) +) # (in-chain: a win at M <= 2, not 4..8) +if SHARED: + assert SH_I in (384, 768, 1536), SH_I + assert SH_PRETILE or SH_SPLIT == 2, ( + "the un-pretiled shared stream is the 32-row (TP4) layout only" + ) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # capacity: fixed: pipeline stages (6 stages per unit) +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) +# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) +KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 + + +def _nkb2(t2: int) -> int: + """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" + return NUM_KB if t2 < KT2 - 1 else KB2_LAST + + +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy +# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few +# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. +HB_REPL = _cfg("hb_repl", 1) +assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) +HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +# Hang guard: every cross-CTA spin (the kernel needs all its CTAs resident at once) traps after this long, so a lost +# co-residency or a missing producer surfaces as a device exception instead of a silent hang. The timer is read only +# inside a spin iteration; the arithmetic is unchanged. +HANG_NS = _cfg("hang_ms", 2000) * 1000000 +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # capacity: fixed: a warp role index +S_WARP = 6 # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +# (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path). +E0_WARP = MMA_WARP # capacity: fixed: a warp role index +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and NRW0 <= 3 and bool(_cfg("e0_off", 1)) +# B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing +# warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the +# two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. +if B1_AFTER_TOPK and XSTAGE and EARLY0 and not E0_OFFLOAD: + raise ValueError( + "m3_moe: b1_after_topk=1 with xstage, early0 and no early-stage offload deadlocks (warp 4 waits on x_bar[1] " + "before TOPK2_BAR, the B1 build waits on TOPK2_BAR before x_bar[1]); set b1_after_topk=0 or early0=0" + ) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 # capacity: fixed: the SM's tensor-memory columns +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# TP2 at 8 tokens: the 98 KB of hidden states fit no 5-stage ring; the router GEMV then reads them from global (the +# kernel's non-XSTAGE path, same math) +if XSTAGE and M_MAX * H * 2 > STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA): + XSTAGE = False +assert not XSTAGE or M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 +_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 +_SH_FIX = ( + SH_GT * SH_NB * 128 + + SH_DKT * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + _SH_XCH_BYTES + + 2 * SH_ROWS * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * SH_NB * 128 +_SB_XC = _SB_BH + SH_DKT * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + _SH_XCH_BYTES +_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = SH_NB +SH_SFA = SH_ACC_DN + 16 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +# pre-tiled stages beyond the ring that warp 0 prefetches into L2, per lane (TP2: 72 stages) +SH_PF = max(2, (SH_STAGES - SH_NST + 31) // 32) if SHARED else 2 +# TP2 (SH_GT 48 gate_up k-tiles, SH_DKT 12) needs more than 256 columns: allocate the next power of two +SH_TMEM = 256 if SH_SFB_DN + 4 * SH_DKT <= 256 else 512 +# SH_TAKE helpers hand the shared allocation to the routed body, whose teardown frees TMEM_COLS +assert not SH_TAKE or SH_TMEM == TMEM_COLS, ( + "SH_TAKE needs the shared TMEM allocation to equal the routed one" +) +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 32 * SH_PF, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 4 * SH_DKT <= SH_TMEM, ( + SH_NST, + _SB_END, + BIG, + ) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _trap_after(elapsed, *, loc=None, ip=None): + """trap when elapsed (ns, int64) exceeds HANG_NS.""" + _llvm.inline_asm( + None, + [cutlass.Int64(elapsed).ir_value(loc=loc, ip=ip)], + f"{{ .reg .pred p; setp.gt.s64 p, $0, {HANG_NS}; @p trap; }}", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _hang_tick(t0): + """One spin iteration of a cross-CTA wait: the wait's start (latched on the first call; t0 = 0 before it), after + trapping once the wait has lasted longer than HANG_NS.""" + now = _gtime() + t1 = cutlass.select_(t0 == cutlass.Int64(0), now, t0) + _trap_after(now - t1) + return t1 + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + # (a partial last k-tile has no atoms past I_TP) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(SH_PF): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(SH_PF) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SH_NB, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if cutlass.const_expr(SH_SPLIT == 1): + # TP2: one K slice; quadrant qd = rows 32 qd .. (gate positions 0..63, then up 0..63), no partials: + # the owner sum is (0 + acc), then bf16, at _SB_GU word (u SH_ROWS + p) 8 + t = (32 qd + l) 8 + t + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + elif cutlass.const_expr(SH_SPLIT == 2): + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + else: + # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are + # accumulator columns 8 qd + t + av = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_GU + qd * cutlass.Int32(8), + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(1): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd == cutlass.Int32(0): + # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + for s_ in cutlass.range_constexpr(1, SH_SPLIT): + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + lane * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(SH_ROWS * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + t0h = cutlass.Int64(0) + while pend != cutlass.Int32(0): + t0h = _hang_tick(t0h) + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(SH_DKT): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + # SH_TAKE helpers keep the TMEM and the epoch for the routed body that follows (its teardown releases both) + not_helper = cutlass.Boolean(True) + if cutlass.const_expr(SH_TAKE): + not_helper = sid < cutlass.Int32(SH_HELPER0) + if (warp == 1) & not_helper: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if (tidx == 0) & not_helper: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if cutlass.const_expr(SH_TAKE): + is_routed = (bx < cutlass.Int32(NCTA)) | (bx >= cutlass.Int32(NCTA + SH_HELPER0)) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if (warp == MMA_WARP) & ( + bx < cutlass.Int32(NCTA) + ): # SH_TAKE helpers reuse the shared allocation + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(PLAN_COLS) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(LATE_TRIGGER and TAIL_TRIGGER and NHEAVY > 0): + late = bx < cutlass.Int32(NHEAVY) + if cutlass.const_expr(SH_TAKE): + late = bx >= cutlass.Int32(NCTA + SH_HELPER0) + if ar_off & late: + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(EARLY_AR): + if ar_on: + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + t0h = cutlass.Int64(0) + while done == cutlass.Int32(0): + t0h = _hang_tick(t0h) + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + t0h = cutlass.Int64(0) + while (cv - cnt_target) < cutlass.Int32(0): + t0h = _hang_tick(t0h) + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + for rr in cutlass.range_constexpr(HB_REPL): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4) + + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + hb_off = cutlass.Int32(0) + if cutlass.const_expr(HB_REPL > 1): + # this CTA's copy of the records + hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + t0h = cutlass.Int64(0) + while done == cutlass.Int32(0): + t0h = _hang_tick(t0h) + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview( + (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) + ), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + t0h = cutlass.Int64(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + t0h = _hang_tick(t0h) + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if cutlass.const_expr(HB_REPL == 1): + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + else: + if lane < HB_REPL: # copy `lane` + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4) + + lane * cutlass.Int32(HB_STRIDE), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + t0h = cutlass.Int64(0) + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + t0h = _hang_tick(t0h) + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/op.py new file mode 100644 index 000000000000..067538849ac8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/op.py @@ -0,0 +1,808 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP2 rank (routed and shared intermediate 1536 per rank), +over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +INTER = 1536 # capacity: fixed: the routed experts' width per tp2 rank (the routed intermediate per TP rank (TP2)) +NUM_EXPERTS = 128 +SHARED_INTER = 1536 # capacity: fixed: the shared expert's width per tp2 rank +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size (m3_moe's shared-expert CTAs (beside num_ctas routed CTAs)) +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +U1 = 2 * INTER // 32 # FC1 units per expert +# The router GEMV takes at most 2 experts per CTA (router weights held in registers). +MIN_CTAS = 64 # capacity: fixed: an SM split, not a size + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # capacity: fixed: pipeline stages (FC1 stages (items) per unit) +_ITEM_BYTES = _UNIT_BYTES // _S1 # capacity: fixed: a ring item's bytes +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters (4, 7 and 8 tokens; at TP2 4 +# tokens with 6 stages measured 0.67 us/layer faster in the ablation chain than the tagged records). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 1, 0, 0, 1, 1) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit; at TP2 9 / 8 / 5 for 1-2 / 3-4 / 8 tokens). 6 at 4 +# tokens measured 1.9 us/layer faster in the ablation chain than the default 8 (same MMAs, bit-identical outputs). +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 6, 0, 0, 0, 0) +# ``early0``: warp 4 streams group 0 right after its own top-k. Off at 8 tokens (0.6 us/layer faster in the chain: +# there it is on every routing warp's critical path and the ring holds only 5 stages). +EARLY0_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 0) +# ``xw``: FC2 B builder warps per token count (0 = the kernel default). 6 at 3 / 4 / 5 / 8 tokens: -1.04 / -0.63 / +# -0.26 / -0.61 us/call (795995, 795401, bit-identical); at 1 / 2 / 6 / 7 tokens +0.87 / -0.07 / +2.79 / +0.05, so not +# there. +XW_BY_M: Tuple[int, ...] = (0, 0, 6, 6, 6, 0, 0, 6) +# ``sh_take``: the routed CTAs owning a second FC2 output block (H / 32 mod num_ctas of them) hand it to as many +# shared-expert CTAs, which run it FC2-only after their shared work; every block keeps its tiles, scales and finalize, +# so the output is bit-identical. On at 6 / 7 / 8 tokens (-1.34 / -4.47 / -1.24 us/call); off at 1 .. 5, where the +# shared work ends too late for the helper (+0.5 .. +8.1). +# mega1.8: OFF by default. Engine-level gates are bitwise vs the kill switch, but the in-server same-mount check +# (797173) read TEXT-SAME 0/7 against a 4/7 base, cause unknown. The table that was gated is (0, 0, 0, 0, 0, 1, 1, 1). +SH_TAKE_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 0) +# The plan's ``heavy_fc2_only``: 192 FC2 output blocks over 188 routed CTAs leave 4 CTAs with two blocks, whose +# trailing FC2 after the last hand-off otherwise ends the kernel 3-7 us after the rest; they stream no FC1 instead +# (1.7 / 3.4 us/layer faster at 4 / 8 tokens in the ablation chain; bit-identical outputs). +HEAVY_FC2_ONLY_BY_M: Tuple[int, ...] = (0, 0, 0, 1, 0, 0, 0, 1) +# ``hb_repl``: copies of the tagged FC1 -> FC2 records (applied when > 1); CTA bx polls copy bx % hb_repl, which +# spreads the polls of the few hot record lines over L2. 7 and 8 tokens hand off through the B images instead. +HB_REPL_BY_M: Tuple[int, ...] = (32, 32, 8, 8, 1, 1, 1, 1) +# ``lg_repl``: replicas of the router logit exchange at up to LG_REPL_MAX_TOKENS tokens (fewer stores before the +# exchange; bit-identical outputs). Above that the kernel's default applies. +LG_REPL = 32 +LG_REPL_MAX_TOKENS = 4 # capacity: fixed: measured crossover +# The TP2 FLAT stack (MoE lane, 2026-10-17): at every token count, FC2 B through per-(slot, token) records (no B2 +# images), 16 router-logit replicas, the AR push parities loaded in the finalize (no preload) and the PDL release right +# after the grid wait. Bitwise-identical output (gate a 791026, mega1 image, m 1..8 x 4 routings); in-model TP2 pair +# -114 us/step. +# The values are part of the engine's module key. +FLAT_STACK = {"b2img": 0, "lg_repl": 16, "ar_pre": 0, "late_trigger": 0} +# FLAT only up to 4 tokens (MoE lane 10-18, 795150): at 8 tokens the per-m tables (B images, preloaded AR push) are +# 2.96 us/call +# faster than FLAT (47.99 vs 50.95; 4 tokens: FLAT -0.51); batch-2 GEN steps run 8 tokens. Bit-identical either way. +FLAT_MAX_TOKENS = 4 # capacity: fixed: measured FLAT crossover + + +def flat_applies(num_tokens: int) -> bool: + """Whether an engine state for ``num_tokens`` tokens runs the FLAT stack (up to FLAT_MAX_TOKENS).""" + return num_tokens <= FLAT_MAX_TOKENS + + +# Byte-balanced FC1 ranges below one unit per CTA (:func:`_fine_ranges`) for groups too small to give every CTA a +# whole unit; False gives such groups whole units on the CTAs with the least FC2 work. +FINE_PLAN = True +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 + = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row + R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then + down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per + layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + P = I_ // n_cta + split = 128 // (2 * P) + gt = H_ // 128 // split + dkt = I_ // 128 + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + octant = r // P + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(2 * dkt): + jt, kt = gd // dkt, gd % dkt + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, gt + gd, :16384] = img.reshape(-1) + out[s, gt + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def _fine_ranges(weights, total: int): + """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item + counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary + strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit + runs at least to that unit's end. Returns the item counts.""" + n = len(weights) + items = [0] * n + s = 0 + acc = 0.0 + for c in range(n): + acc += weights[c] + e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) + if s % _S1 and e < total: # started inside a unit: run to (at least) its end + e = max(e, min(total, (s // _S1 + 1) * _S1)) + items[c] = e - s + s = e + assert s == total + return items + + +def make_phase_plan( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = FINE_PLAN, + heavy_fc2_only: bool = False, +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + # (heavy_fc2_only) CTAs with more FC2 blocks than the minimum stream no FC1: they run each group's FC2 as soon as + # its hand-off lands, so their trailing FC2 after the last hand-off matches the others' two lagged groups + min_nb = min(nbs) + no_fc1 = [heavy_fc2_only and nbs[c] > min_nb for c in range(num_ctas)] + BIG_W = 1 << 40 + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + for c in range(num_ctas): + if no_fc1[c]: + fc2[c] += BIG_W + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [ + 0 if no_fc1[c] else max(_S1, (limit - f) // _ITEM_BYTES) + for c, f in enumerate(fc2) + ] + + lo, hi = 0, max(f for c, f in enumerate(fc2) if not no_fc1[c]) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 + lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] + sw = sum(wts) + items = _fine_ranges([w * total / sw for w in wts], total) + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +def _shared_takes_fc2_plan(plan: torch.Tensor, num_ctas: int, m_max: int) -> torch.Tensor: + """The plan with SHARED_CTAS more columns for SH_TAKE: routed CTA c < H / 32 mod num_ctas keeps only its first FC2 + block (its FC1 rows unchanged); helper shared CTA SHARED_CTAS - nheavy + c owns the second, with no FC1.""" + nheavy = (HIDDEN // 32) % num_ctas + ext = torch.zeros(plan.shape[0], num_ctas + SHARED_CTAS, plan.shape[2], dtype=plan.dtype) + ext[:, :num_ctas] = plan + h0 = num_ctas + SHARED_CTAS - nheavy + for c in range(nheavy): + b0 = int(plan[TOP_K, c, m_max]) & 0xFFFF + assert (int(plan[TOP_K, c, m_max]) >> 16) == 2 + ext[:, c, m_max] = b0 | (1 << 16) + ext[:, h0 + c, m_max] = (b0 + 1) | (1 << 16) + return ext + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + inter: int = INTER, + early_release: bool = True, + ): + """``inter``: the routed intermediate per rank; this engine is built for INTER (1536) only. + ``early_release``: on a call that pushes into the boundary all-reduce (``arp``), every CTA releases the PDL + dependents right after its grid dependency wait (the kernel's EARLY_AR); False releases them after the push.""" + if inter != INTER: + raise ValueError( + f"m3_moe (TP2): routed intermediate {inter} per rank is not supported (only {INTER})" + ) + self.inter = inter + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + if self.num_ctas > HIDDEN // 32: + # a routed CTA that owns no 32-row FC2 output block deadlocks the FC2 hand-off on its second call + raise ValueError( + f"m3_moe runs at most {HIDDEN // 32} routed CTAs (one FC2 output block each), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, + num_tokens, + M3_MOE_LAG, + 0.0 if num_tokens == 1 else 1.0, + True, + heavy_fc2_only=bool(HEAVY_FC2_ONLY_BY_M[num_tokens - 1]), + ) + sh_take = ( + shared + and SH_TAKE_BY_M[num_tokens - 1] > 0 + and self.num_ctas < HIDDEN // 32 + and 0 < (HIDDEN // 32) % self.num_ctas <= SHARED_CTAS + ) + if sh_take: + plan = _shared_takes_fc2_plan(plan, self.num_ctas, num_tokens) + from tensorrt_llm.logger import logger as _shlog + + _shlog.info( + f"MiniMax-M3 MoE engine SH_TAKE (TP2): {(HIDDEN // 32) % self.num_ctas} shared CTAs take the extra " + f"FC2 blocks, tokens={num_tokens}" + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + **({"sh_take": 1} if sh_take else {}), + "timing": 0, + "i_tp": INTER, + } + if HB_REPL_BY_M[num_tokens - 1] > 1: + self.config["hb_repl"] = HB_REPL_BY_M[num_tokens - 1] + if num_tokens <= LG_REPL_MAX_TOKENS: + self.config["lg_repl"] = LG_REPL + if shared: + self.config["sh_i"] = SHARED_INTER + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if not early_release: + self.config["early_ar"] = 0 + if XW_BY_M[num_tokens - 1] > 0: + self.config["xw"] = XW_BY_M[num_tokens - 1] + from tensorrt_llm.logger import logger as _xlog + + _xlog.info( + f"MiniMax-M3 MoE engine XW (TP2): xw={XW_BY_M[num_tokens - 1]} tokens={num_tokens}" + ) + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + if not EARLY0_BY_M[num_tokens - 1]: + self.config["early0"] = 0 + if flat_applies(num_tokens): + self.config.update(FLAT_STACK) + # engagement line (every rank; once per engine state = once per token count) + from tensorrt_llm.logger import logger as _flog + + _flog.info( + "MiniMax-M3 MoE engine FLAT: " + + " ".join(f"{k}={self.config[k]}" for k in FLAT_STACK) + + f" tokens={num_tokens}" + ) + else: + from tensorrt_llm.logger import logger as _flog + + _flog.info( + f"MiniMax-M3 MoE engine FLAT off above {FLAT_MAX_TOKENS} tokens (per-m tables): tokens={num_tokens}" + ) + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + # Every CTA waits on others (cross-CTA spins on partials, hand-off records, the shared expert's h and y): the + # grid must be resident at once, one CTA per SM. The module's shared-memory / TMEM budget assertion makes one + # CTA fit an SM; the grid must not exceed the SMs. (The kernel also traps after HANG_NS in any such spin.) + if self.num_ctas + mod.NSH > sms: + raise ValueError( + f"m3_moe: a grid of {self.num_ctas} + {mod.NSH} CTAs exceeds the {sms} SMs; its CTAs wait on each " + "other and must all be resident at once" + ) + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full( + (mod.G_CAP * mod.U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw + ) + # the FC1 -> FC2 records; with HB_REPL > 1 copies, HB_STRIDE int32 words apart + hb_words = mod.G_CAP * num_tokens * mod.U1 * 4 + if mod.HB_REPL > 1: + hb_words = mod.HB_REPL * mod.HB_STRIDE + self.hbuf32 = torch.zeros(hb_words, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [3072, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 1536], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + # keyed on the config the kernel module was built with (not self.config), so a module built with any other + # trace-time config can never reuse another build's compiled kernel + key = (_ENGINE_PATH,) + tuple(sorted(self.mod.M3_CONFIG.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe_finalize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe_finalize.py new file mode 100644 index 000000000000..714eafc4661b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe_finalize.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exact MoE finalize for the routed experts' unfinalized output, optionally with the shared expert added. + +The TRTLLM-Gen MoE runner with ``do_finalize=False`` returns the FC2 output in permuted order, the per-token top-k +routing weights and ``expanded_idx_to_permuted_idx``. Its own finalize (``finalizeKernelVecLoad``) computes, per token +and hidden element, ``acc = 0; acc += weight_k * fc2[permuted_idx_k]`` for k in order (skipping ``-1``) in float32 +and rounds to BF16 once, but issues the k gathers one after another, so at prefill sizes it runs ~2.6x its byte +floor (153 us per layer at 16K rows on SM107). This kernel issues the gathers together and keeps the per-element +order, so the result is bitwise the runner's: a BF16 x BF16 product is exact in float32, so fused and unfused +multiply-add agree as well. ``shared`` adds the shared expert as ``shared_output.add_(routed_output)`` does +(``bf16(float(shared) + float(routed))``), which removes that separate pass over the hidden states. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +#: (BLOCK_H, num_warps): the tile only splits the hidden dimension, so it never changes the bits. Tuned on SM107 +#: (2K / 4K / 16K rows: 11.2 / 20.5 / 74.9 us, 1.52 / 1.39 / 1.27x the byte floor; the runner's finalize +#: 23 / 40 / 153 us). +CONFIGS = ((2048, 4),) + + +@triton.jit +def _moe_finalize_kernel(fc2, idx, weights, shared, out, ld_fc2, H: tl.constexpr, TOPK: tl.constexpr, + BLOCK_H: tl.constexpr, SHARED: tl.constexpr): # fmt: skip + t = tl.program_id(0).to(tl.int64) + cols = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + cmask = cols < H + acc = tl.zeros((BLOCK_H,), dtype=tl.float32) + for k in tl.static_range(TOPK): + p = tl.load(idx + t * TOPK + k) + w = tl.load(weights + t * TOPK + k).to(tl.float32) + valid = p >= 0 + x = tl.load( + fc2 + tl.maximum(p, 0).to(tl.int64) * ld_fc2 + cols, mask=cmask & valid, other=0.0 + ) + acc = tl.where(valid, tl.fma(w, x.to(tl.float32), acc), acc) + r = acc.to(tl.bfloat16) + if SHARED: + s = tl.load(shared + t * H + cols, mask=cmask, other=0.0) + r = (s.to(tl.float32) + r.to(tl.float32)).to(tl.bfloat16) + tl.store(out + t * H + cols, r, mask=cmask) + + +def m3_moe_finalize( + fc2: torch.Tensor, + weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + shared: Optional[torch.Tensor] = None, +) -> Optional[torch.Tensor]: + """BF16 [tokens, H] = sum_k weights[:, k] * fc2[expanded_idx_to_permuted_idx[:, k]] (+ ``shared``), bitwise to the + TRTLLM-Gen finalize (and to ``shared.add_(that)``); None when the operands are not the runner's layout.""" + if ( + fc2.dim() != 2 + or fc2.dtype != torch.bfloat16 + or fc2.stride(1) != 1 + or weights.dim() != 2 + or expanded_idx_to_permuted_idx.shape != weights.shape + or expanded_idx_to_permuted_idx.dtype != torch.int32 + or not weights.is_contiguous() + or not expanded_idx_to_permuted_idx.is_contiguous() + or ( + shared is not None + and ( + shared.dtype != torch.bfloat16 + or not shared.is_contiguous() + or shared.shape != (weights.shape[0], fc2.shape[1]) + ) + ) + ): + return None + tokens, topk = weights.shape + hidden = fc2.shape[1] + out = torch.empty((tokens, hidden), dtype=torch.bfloat16, device=fc2.device) + if tokens == 0: + return out + block_h, warps = CONFIGS[0] + _moe_finalize_kernel[(tokens, triton.cdiv(hidden, block_h))]( + fc2, + expanded_idx_to_permuted_idx, + weights, + shared if shared is not None else out, + out, + fc2.stride(0), + hidden, + topk, + block_h, + shared is not None, + num_warps=warps, + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107__tp2/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_107__tp2/kernels/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_report.py new file mode 100644 index 000000000000..7f864bd240c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_report.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast-path reports a kernel module makes from inside a call: ``taken`` (an engagement line), ``missed`` (a covered +shape that cannot run: raises under ``modeling_v2='require'``), ``uncovered`` and ``capture_missed``. + +The policy (logging once, raising under require) is the target's: ``modeling.py`` installs its functions here at import, +so a kernel reports without importing the target module. Until then every report is a no-op. +""" + +from __future__ import annotations + +from typing import Callable, Optional + + +def _noop(*args, **kwargs) -> None: + return None + + +taken: Callable[..., None] = _noop +missed: Callable[..., None] = _noop +uncovered: Callable[..., None] = _noop +capture_missed: Callable[..., None] = _noop + + +def install( + *, + on_taken: Callable[[str, Optional[int], str], None], + on_missed: Callable[[str, Optional[int], str], None], + on_uncovered: Callable[[str, str], None], + on_capture_missed: Callable[[str, Optional[int], str], None], +) -> None: + """Route this module's reports to the target's policy.""" + global taken, missed, uncovered, capture_missed + taken, missed, uncovered, capture_missed = on_taken, on_missed, on_uncovered, on_capture_missed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_router_gemv.py new file mode 100644 index 000000000000..30855eb7f05d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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 # capacity: fixed: the device's shared-memory limit + + +def m3_router_gemv( + x: torch.Tensor, weight: torch.Tensor, rows: int = 1, use_pdl: bool = False +) -> Optional[torch.Tensor]: + """fp32 ``x @ weight.T`` for bf16 ``x`` with at most 16 rows; None when unsupported (only ``rows=1`` is built).""" + if ( + rows != 1 + or not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or weight.shape[0] % rows != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or rows * x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return _router_gemv_op(x, weight, rows, use_pdl) + + +def m3_router_gemv_fp4( + x: torch.Tensor, weight: torch.Tensor, sf_scale: torch.Tensor, use_pdl: bool = False +) -> Optional[tuple]: + """Router logits plus the NVFP4 activation the routed experts consume (bit-identical to + ``fp4_quantize(x, sf_scale, 16, False, False)``): (logits [M, E] fp32, x_fp4 [M, K / 2] uint8, + x_sf [M * K / 16] uint8 linear). None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or sf_scale is None + or sf_scale.dtype != torch.float32 + or sf_scale.numel() != 1 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return tuple(_router_gemv_fp4_op(x, weight, sf_scale.reshape(1), use_pdl)) + + +def m3_router_tc( + x: torch.Tensor, + weight: torch.Tensor, + sf_scale: Optional[torch.Tensor] = None, + use_pdl: bool = False, +) -> Optional[tuple]: + """Router logits (TF32 tensor cores, like the native TF32 GEMM) for 1..48 tokens with 128 experts, plus the + NVFP4 activation when ``sf_scale`` is given. Returns (logits,) or (logits, x_fp4, x_sf); None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 48 + or weight.dtype != torch.float32 + or weight.shape[0] != 128 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or (x.stride(0) * 2) % 16 != 0 + or x.data_ptr() % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or 16 * (x.shape[1] // 8 + 4) * 4 + x.shape[0] * (x.shape[1] // 8 + 8) * 2 + > _MAX_DYNAMIC_SMEM + or (sf_scale is not None and (sf_scale.dtype != torch.float32 or sf_scale.numel() != 1)) + ): + return None + res = _router_tc_op(x, weight, None if sf_scale is None else sf_scale.reshape(1), use_pdl) + return tuple(res) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_sampled_accept.py new file mode 100644 index 000000000000..93ffda80ecea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_sampled_accept.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). + +Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +active must stay on the generic path; the caller decides. + +Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +(the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens heavier +than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the nucleus +with probability >= top_p, so most rows decide on their second exchange. + +Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based Philox +keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one 4-word +record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and reads +the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the same +order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing and +two call halves keep a rank one call ahead from overwriting a slot still being read. +""" + +from typing import Optional + +import torch + +from . import m3_report + +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu) and called as +# torch.ops.trtllm.minimax_m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) +# The TP2 vocab shard limit of trtllm::minimax_m3_sm107tp2_sampled_accept (kMaxVs): 200064 / 2 = 100032 columns. +MAX_VS = 100352 # capacity: fixed: the vocabulary shard bound at tp2 + + +class SampledAcceptWorkspace: + """Exchange slots behind one multicast mapping (two call halves), the (finished rows, call) counters, the per-row + token scratch and an extra-attempt counter. Created collectively, outside CUDA-graph capture.""" + + def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): + self.uc, self.mc_ptr, self.rank, self.world = uc, mc_ptr, rank, world + self.handle, self.comm = handle, comm + dev = uc.device + self.counters = torch.zeros(2, dtype=torch.int32, device=dev) + self.tokens = torch.zeros(MAX_ROWS, dtype=torch.int32, device=dev) + self.extra_attempts = torch.zeros(1, dtype=torch.int32, device=dev) + + @classmethod + def create(cls, mapping, device) -> "SampledAcceptWorkspace": + words = int( + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept_ws_words(max(mapping.tp_size, 1)) + ) + if mapping.tp_size == 1: + uc = torch.zeros(words, dtype=torch.int64, device=device) + return cls(uc, uc.data_ptr(), 0, 1) + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 8, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("SampledAcceptWorkspace.create failed on another rank") + return cls(uc, mc.data_ptr(), mapping.tp_rank, mapping.tp_size, handle, comm) + + +_WS: Optional[SampledAcceptWorkspace] = None + + +def workspace(mapping, device) -> Optional[SampledAcceptWorkspace]: + global _WS + if _WS is None: + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("spec_tail", None, "the sampled-accept workspace") + return None + _WS = SampledAcceptWorkspace.create(mapping, device) + return _WS + + +def sampled_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: Optional[torch.Tensor], + seeds: torch.Tensor, + offsets: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, + ws: Optional[SampledAcceptWorkspace] = None, + use_pdl: bool = True, +): + """Sampled verification (temperature + top-p) from this rank's bf16 logits shard [nreq * (draft_len + 1), V / tp]: + outputs as shard_accept (accepted [nreq, D+1], num_accepted [nreq], gather_ids [nreq]); None + when unavailable. Per-row sampling parameters and Philox (seed, offset) are the spec metadata's, one per row. + top_ks / min_ps are not applied: rows with top-k or min-p active must be routed to the generic path by the + caller.""" + del top_ks, min_ps + nreq, draft_len = draft_tokens.shape + rows = nreq * (draft_len + 1) + if logits.dim() != 2 or logits.shape[0] != rows or rows > MAX_ROWS or logits.shape[1] > MAX_VS: + return None + if logits.dtype != torch.bfloat16 or logits.stride(1) != 1: + return None + ws = ws if ws is not None else workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + temperatures = temperatures[:rows].float().contiguous() + top_ps = top_ps[:rows].float().contiguous() + seeds = seeds[:rows].long().contiguous() + offsets = offsets[:rows].long().contiguous() + _launch( + logits, vocab_offset, temperatures, top_ps, seeds, offsets, ws, draft_tokens, batch_idx, + accepted, num_accepted, gather_ids, force_total, next_new, force_vec, use_pdl, rows, + ) # fmt: skip + return accepted, num_accepted, gather_ids + + +def _launch( + logits, vocab_offset, temperatures, top_ps, seeds, offsets, ws, draft_tokens, batch_idx, + accepted, num_accepted, gather_ids, force_total, next_new, force_vec, use_pdl, rows, +) -> None: # fmt: skip + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept( + logits, + int(vocab_offset), + temperatures, + top_ps, + seeds, + offsets, + ws.uc, + ws.mc_ptr, + ws.counters, + ws.tokens, + ws.rank, + ws.world, + draft_tokens.int().contiguous(), + batch_idx.int().contiguous(), + accepted, + num_accepted, + gather_ids, + int(force_total), + next_new, + force_vec, + ws.extra_attempts, + bool(use_pdl), + PREWAIT_PARAMS, + EARLY_TRIGGER, + _ilp_walk(rows), + ) + + +# The sampled verify loads each row's sampling parameters before its PDL dependency wait (overlapping the head GEMM's +# tail; the same values and tokens), and releases its dependents right after that wait (the dependent parks at its +# own wait; the same values). +PREWAIT_PARAMS = True +EARLY_TRIGGER = True +# SA-PW: the owner walk takes its tokens 8 at a time (the same draw, tokens and Philox draws as the one-token walk, +# bit for bit; -2.2..-2.5 us per call on Rubin TP2). Off in the TP2 deployment (frozen deploy value). +ILP_WALK = False +ILP_WALK_PATH = "m3_sampled_accept_ilp_walk" + + +def _ilp_walk(rows: int) -> bool: + """ILP_WALK, logged once per row count as taken; off, one 'not taken' line (logged once).""" + if not ILP_WALK: + m3_report.uncovered(ILP_WALK_PATH, "the one-token owner walk (ILP_WALK off)") + return False + m3_report.taken( + ILP_WALK_PATH, + rows, + "the sampled verify's owner walk takes 8 tokens at a time", + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_select_long_dsl.py new file mode 100644 index 000000000000..6e6e23e0c0b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_select_long_dsl.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). + +Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long``, +one 1024-thread CTA per row), from a CTA of 384 threads: + +* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and + mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's + valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned + key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped + block and masks the result, so a thread's loads issue back to back. +* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of + the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. +* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by + that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate + against the whole row instead. +* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. + +The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait +(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are +compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and +loads them before the executor's warmup, so a decode step never compiles one. +""" + +from typing import Dict, Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.utils import ( + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +# The 384-thread select reads its row's valid-block count before its grid wait: off (its TP2 in-model pairs did not +# pass); the selection is the same either way. +SELECT_N_VALID_PREWAIT = False + +SELECT_THREADS = 384 +SELECT_WARPS = SELECT_THREADS // 32 +# 1M tokens + draft tokens at 128 tokens per block. +MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row +MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging +TOPK = 16 # capacity: fixed: MSA's top-k +NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) +# The key of -inf (the reference's 0xff800000): never emitted. +NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 +# (keys per thread, rows up to this many valid blocks) +ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) + +# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare +# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group +# minima [16], per-warp candidate counts [SELECT_WARPS]. +_CAND = 0 +_PICK_KEY = _CAND + 2 * MAX_CAND +_PICK_BLOCK = _PICK_KEY + TOPK +_WARP_MIN = _PICK_BLOCK + TOPK +_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 +_WARP_CNT = _GROUP_MIN + 16 +SMEM_WORDS = _WARP_CNT + SELECT_WARPS +SMEM_BYTES = 4 * SMEM_WORDS + + +def _key(score): + """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" + bits = cutlass.Float32(score).bitcast(cutlass.Int32) + return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) + + +def _imin(a, b): + return cutlass.select_(b < a, b, a) + + +def _imax(a, b): + return cutlass.select_(b > a, b, a) + + +def _forced_key(scores, head, block, query, init_blocks, local_start): + raw = scores[head, block, query] + s = cutlass.select_( + block >= local_start, + cutlass.Float32(LOCAL_SCORE), + cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), + ) + return _key(s) + + +@cute.jit +def _select_row( + row, + scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides + n_valid: cute.Tensor, # int32 [>= total_q] + out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] + init_blocks, + local_blocks, + num_kv_heads, + smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory + bar_id: cutlass.Constexpr, + loaded_trigger: cutlass.Constexpr, + raw_prewait, + n_valid_prewait: cutlass.Constexpr, +): + """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier + ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the + row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by + the caller before its grid wait; otherwise it is read here.""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % cutlass.Int32(32) + warp = tid // cutlass.Int32(32) + query = row // num_kv_heads + head = row % num_kv_heads + if cutlass.const_expr(n_valid_prewait): + raw = raw_prewait + else: + raw = n_valid[query] + valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) + local_start = _imax(raw - local_blocks, cutlass.Int32(0)) + cand = cute.make_tensor( + cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) + ) + + if tid < TOPK: + smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) + smem[_PICK_BLOCK + tid] = cutlass.Int32(0) + + if valid <= TOPK: + # Every valid block is picked (thread t holds block t). + if tid < valid: + smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) + smem[_PICK_BLOCK + tid] = tid + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): + items = ITEM_CLASSES[ci][0] + lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] + hi = ITEM_CLASSES[ci][1] + if (valid > lo) & (valid <= hi): + key = cute.make_rmem_tensor(items, cutlass.Int32) + local_min = cutlass.Int32(NO_KEY) + for s in cutlass.range_constexpr(items): + # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. + block = tid + cutlass.Int32(s * SELECT_THREADS) + k = _forced_key( + scores, head, _imin(block, valid - 1), query, init_blocks, local_start + ) + key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) + local_min = _imin(local_min, key[s]) + local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) + if lane < 16: + smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + if tid < 16: + g = cutlass.Int32(NO_KEY) + for w in cutlass.range_constexpr(SELECT_WARPS): + g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) + smem[_GROUP_MIN + tid] = g + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + bound = cutlass.Int32( + -2147483648 + ) # the smallest key: signed keys, unlike the reference's unsigned ones + for g in cutlass.range_constexpr(16): + bound = _imax(bound, smem[_GROUP_MIN + g]) + + # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(items): + cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) + incl = cnt + for d in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync( + incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) + ) + incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) + if lane == 31: + smem[_WARP_CNT + warp] = incl + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + before = cutlass.Int32(0) + total = cutlass.Int32(0) + for w in cutlass.range_constexpr(SELECT_WARPS): + c = smem[_WARP_CNT + w] + before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) + total += c + slot = before + incl - cnt + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + if slot < MAX_CAND: + cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( + tid + cutlass.Int32(s * SELECT_THREADS) + ) + slot += 1 + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if total <= MAX_CAND: + for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): + mine = cand[i] + rank = cutlass.Int32(0) + for j in cutlass.range(total, unroll=8): + rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) + smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) + else: + # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + block = tid + cutlass.Int32(s * SELECT_THREADS) + rank = cutlass.Int32(0) + for j in cutlass.range(valid, unroll=1): + kj = _forced_key(scores, head, j, query, init_blocks, local_start) + less = (kj < key[s]) | ((kj == key[s]) & (j < block)) + rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = key[s] + smem[_PICK_BLOCK + rank] = block + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if warp == 0: + # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. + v = cutlass.Int32(NO_KEY) + if lane < TOPK: + pk = smem[_PICK_KEY + lane] + if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): + v = smem[_PICK_BLOCK + lane] + for size_log in cutlass.range_constexpr(1, 6): + for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): + other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) + ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) + lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) + v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) + if lane < TOPK: + out[query, head, lane] = cutlass.select_( + v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v + ) + + +class _SelectLongDslKernel: + """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count + before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because + ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's + first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it + (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" + + def __init__( + self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool + ): + self.use_pdl = use_pdl + self.early_trigger = early_trigger + self.loaded_trigger = (not early_trigger) and trigger_phase == 1 + self.n_valid_prewait = n_valid_prewait + + @cute.kernel + def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): + row, _, _ = cute.arch.block_idx() + smem = SmemAllocator() + words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) + raw_prewait = cutlass.Int32(0) + if cutlass.const_expr(self.n_valid_prewait): + raw_prewait = n_valid[row // num_kv_heads] + griddepcontrol_wait() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + _select_row( + row, + scores, + n_valid, + out, + init_blocks, + local_blocks, + num_kv_heads, + words, + bar_id=1, + loaded_trigger=self.loaded_trigger, + raw_prewait=raw_prewait, + n_valid_prewait=self.n_valid_prewait, + ) + + @cute.jit + def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): + num_kv_heads = scores.shape[0] + self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( + grid=(scores.shape[2] * num_kv_heads, 1, 1), + block=(SELECT_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + +#: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). +kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} + + +def kernel_key( + use_pdl: bool, early_trigger: bool, trigger_phase: int +) -> Tuple[bool, bool, int, bool]: + """The compile key of a launch (the trigger phase does not matter with the early trigger), with this target's + ``SELECT_N_VALID_PREWAIT``.""" + return ( + bool(use_pdl), + bool(early_trigger), + 0 if early_trigger else int(trigger_phase), + bool(SELECT_N_VALID_PREWAIT), + ) + + +def _compile(key: Tuple[bool, bool, int, bool]): + if key not in kernel_cache: + heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() + scores = make_fake_tensor( + cutlass.Float32, + (heads, blocks, rows), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) + out = make_fake_tensor( + cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) + ) + kernel_cache[key] = cute.compile( + _SelectLongDslKernel(*key), + scores, + n_valid, + out, + cutlass.Int32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + return kernel_cache[key] + + +def select_long_dsl( + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool, + early_trigger: bool, + trigger_phase: int, +) -> None: + """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" + _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( + scores, n_valid_blocks, out, init_blocks, local_blocks + ) + + +def prebuild(keys, device) -> list: + """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" + scores = torch.zeros( + 1, 2048, 4, dtype=torch.float32, device=device + ) # capacity: fixed: a dummy prebuild row, not a size + n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) + out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) + built = [] + for key in keys: + key = kernel_key(*key) + _compile(key)(scores, n_valid, out, 0, 1) + built.append(key) + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_shard_accept.py new file mode 100644 index 000000000000..8382d02fb031 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_shard_accept.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 verification for the MiniMax-M3 target head (vocab-sharded TP head, verify rows <= 48). + +Replaces the all-gather of the target logits, the argmax and the strict-acceptance ops of a greedy verify step with +one kernel over this rank's BF16 logits shard (the target lm_head's own local GEMM output): + +* Every logit becomes a 64-bit key ``ordered(value) << 32 | ~global_index`` (ties pick the lowest index, as + ``torch.argmax``), max-reduced per row: warp shuffles, then ``atomicMax``. +* The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping + (Lamport: 0 = empty, the key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the + slots and runs the strict acceptance. Call n uses slot half n % 2. + +The argmax is taken over the BF16 logits without arithmetic, so the tokens are those of the in-tree path's argmax +over the gathered logits. The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/ +shardAccept.cu) and called as ``torch.ops.trtllm.minimax_m3_shard_accept``. +""" + +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from . import m3_report + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 # capacity: fixed: the kernel's built verify rows +# A verify step of more than SHARD_ROWS rows (13..48 requests at the built draft length) runs in chunks of whole +# requests, each chunk a call of the certified geometry on its own rows; greedy chunks are bitwise the gathered argmax +# (lowest index on ties). Rubin TP2 W1b 797141. Kill: False (past SHARD_ROWS the gathered logits + argmax run). +WIDE_SPEC_TAIL = True + + +@dataclass(eq=False) +class AcceptWorkspace: + """Per-rank exchange slots ([2 halves][SHARD_ROWS rows][world] uint64) behind one multicast mapping, the local + key accumulators and the (finished CTAs, call) counters.""" + + uc: torch.Tensor + mc_ptr: int + keys: torch.Tensor + counters: torch.Tensor + rank: int + world_size: int + handle: Any = None + comm: Any = None + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + if not _mnnvl_workspace_all_succeeded(comm, True): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = mc.data_ptr(), mapping.tp_rank + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None + return AcceptWorkspace( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) + + +def _workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace, created collectively outside CUDA-graph capture on first use.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("m3_shard_accept", None, "the shard-accept workspace") + return None + _WS = create_workspace(mapping, device, SHARD_ROWS) + return _WS + + +def shard_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, +): + """Greedy verification from this rank's BF16 logits shard. + + Args: + logits: [nreq * (draft_len + 1) <= SHARD_ROWS, V / tp] BF16, unit column stride: this rank's columns + (global ids ``vocab_offset ..``) of the verify rows, request-major. + draft_tokens: [nreq, draft_len] int32, the drafts being verified. + batch_idx: [nreq] int32, each request's row in the verify layout. + force_total / force_vec: forced acceptance (see ``m3_shard_accept``). + next_new: optional int32 [nreq, >= 1]; column 0 receives the last accepted token. + + Returns: + (accepted_tokens [nreq, draft_len + 1] int32, num_accepted [nreq] int32, first-pass gather ids [nreq] + int64), or None when the shape is not served or the workspace is missing during a capture. Collective over + the TP group. + """ + nreq, draft_len = draft_tokens.shape + if ( + logits.dim() != 2 + or logits.shape[0] != nreq * (draft_len + 1) + or logits.dtype != torch.bfloat16 + or logits.stride(1) != 1 + ): + return None + if logits.shape[0] > SHARD_ROWS and not (WIDE_SPEC_TAIL and SHARD_ROWS // (draft_len + 1) >= 1): + return None + ws = _workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + if logits.shape[0] > SHARD_ROWS: + # WIDE_SPEC_TAIL: whole requests per launch (SHARD_ROWS // (D + 1) of them), each chunk a call of the + # certified geometry on its own rows; batch_idx is the identity layout (Eagle3 metadata), so a chunk passes + # its local rows and its gather ids are shifted back to the step's verify layout. + per = SHARD_ROWS // (draft_len + 1) + local = _local_rows(per, dev) + for c0 in range(0, nreq, per): + c1 = min(nreq, c0 + per) + r0, r1 = c0 * (draft_len + 1), c1 * (draft_len + 1) + _launch( + logits[r0:r1], vocab_offset, ws, draft_tokens[c0:c1], local[: c1 - c0], force_total, + None if next_new is None else next_new[c0:c1], + None if force_vec is None else force_vec[c0:c1], + accepted[c0:c1], num_accepted[c0:c1], gather_ids[c0:c1], + ) # fmt: skip + gather_ids[c0:c1] += r0 + m3_report.taken( + "wide_spec_tail", logits.shape[0], f"shard accept in {-(-nreq // per)} chunks" + ) + return accepted, num_accepted, gather_ids + _launch( + logits, vocab_offset, ws, draft_tokens, batch_idx, force_total, next_new, force_vec, + accepted, num_accepted, gather_ids, + ) # fmt: skip + return accepted, num_accepted, gather_ids + + +_LOCAL_ROWS: dict = {} + + +def _local_rows(n: int, device) -> torch.Tensor: + """A cached int32 arange(n): a chunk's batch_idx in its own (identity) verify layout.""" + key = (n, device) + t = _LOCAL_ROWS.get(key) + if t is None: + t = _LOCAL_ROWS[key] = torch.arange(n, dtype=torch.int32, device=device) + return t + + +def _launch( + logits, vocab_offset, ws, draft_tokens, batch_idx, force_total, next_new, force_vec, + accepted, num_accepted, gather_ids, +) -> None: # fmt: skip + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_tokens.contiguous(), + batch_idx, + accepted, + num_accepted, + gather_ids, + force_total, + next_new, + _num_sms(), + force_vec, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107__tp2/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_107__tp2/kernels/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107__tp2/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_107__tp2/kernels/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_tokmap_guard.py new file mode 100644 index 000000000000..db0077117fe3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_tokmap_guard.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The token map of ``thop.trtllm_gen_generation_preprocess`` for one request with several tokens. + +For multi-token generation the op's QKV preprocessing places every row by the workspace token map (``tokensInfo``, +(request, token) per row), but builds that map only when ``isBuildDecoderInfoKernelNeeded()`` holds, which at one +request with the default RoPE it does not. The op then applies whatever map the last multi-request call left in the +workspace: a row can land on another request's length and pages, and its RoPE position and K/V slot are wrong. This +module wraps the op so that, for one request and several tokens, the map the build would write, (0, t) for each row t, +is written into the op's own map slice of the workspace (the offset from the op's layout function) just before the +call. Every other call passes through unchanged. The rows' positions still come from the request's generation length, +exactly as with a freshly built map. + +The write is a device copy from a map prebuilt at install (no allocation, no host sync), so it is captured into the +decode graphs with the call it precedes. +""" + +from __future__ import annotations + +from typing import Callable, Dict, Optional, Tuple + +import torch + +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.logger import logger + +#: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). +_ARG_NAMES = ( + "qkv_input", + "workspace", + "sequence_lengths", + "spec_decoding_generation_lengths", + "spec_decoding_position_offsets", + "kv_cache_block_offsets", + "host_kv_cache_pool_pointers", + "host_kv_cache_pool_mapping", + "kv_scale_orig_quant", + "kv_scale_quant_orig", + "attention_output_orig_quant", + "rotary_inv_freq", + "rotary_cos_sin", + "mrope_position_deltas", + "layer_idx", + "seq_offset", + "num_heads", + "num_kv_heads", + "head_size", + "tokens_per_block", + "kv_cache_quant_mode", + "max_attention_window_size", + "cyclic_attention_window_size", + "num_tokens", + "batch_beam", + "input_seq_length", + "max_past_kv_length", + "rotary_embedding_dim", + "rotary_embedding_base", + "rotary_embedding_scale_type", + "rotary_embedding_scale", + "rotary_embedding_max_positions", + "position_embedding_type", + "bmm1_scale", + "bmm2_scale", + "fp8_context_fmha", + "predicted_tokens_per_seq", + "attention_chunk_size", + "multi_processor_count", + "total_num_blocks", + "kv_factor", + "need_build_kv_cache_metadata", + "cross_attention", + "skip_fmha_workspace", +) +_INT2_BYTES = 8 # capacity: fixed: an element size + +_original: Optional[Callable] = None +_max_rows = 0 +_maps: Dict[torch.device, torch.Tensor] = {} +_offsets: Dict[Tuple, int] = {} + + +def _no_taken(rows: int) -> None: + pass + + +def _no_missed(rows: int, reason: str) -> None: + pass + + +_on_taken: Callable[[int], None] = _no_taken +_on_missed: Callable[[int, str], None] = _no_missed + + +def _indexed(device: torch.device) -> torch.device: + """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: + """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" + device = _indexed(device) + m = _maps.get(device) + if m is None or m.shape[0] < rows: + reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" + _on_missed(rows, reason) + raise RuntimeError(f"m3_tokmap_guard: {reason}") + return m[:rows] + + +def _map_offset(args: dict, rows: int) -> int: + x = args["qkv_input"] + skip = bool(args.get("skip_fmha_workspace", False)) + key = ( + x.dtype, + rows, + int(args["num_heads"]), + int(args["head_size"]), + int(args["rotary_embedding_dim"]), + int(args["num_kv_heads"]), + skip, + ) + off = _offsets.get(key) + if off is None: + layout = thop.get_trtllm_gen_generation_workspace_layout( + x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip + ) + off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) + if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: + raise RuntimeError( + f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " + f"(offset {off}, {size} bytes): the layout changed under this guard" + ) + _offsets[key] = off + return off + + +def _guarded(*args, **kwargs): + a = dict(zip(_ARG_NAMES, args)) + a.update(kwargs) + rows = int(a["num_tokens"]) + if ( + a.get("spec_decoding_generation_lengths") is not None + and int(a["predicted_tokens_per_seq"]) > 1 + and int(a["batch_beam"]) == 1 + and rows > 1 + and not bool(a.get("cross_attention", False)) + ): + ws = a["workspace"] + src = _one_request_map(ws.device, rows) + off = _map_offset(a, rows) + # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. + dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] + dst.view(torch.int32).view(rows, 2).copy_(src) + _on_taken(rows) + return _original(*args, **kwargs) + + +def install( + device: torch.device, + max_rows: int, + on_taken: Callable[[int], None], + on_missed: Callable[[int, str], None] = _no_missed, +) -> bool: + """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to + ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this + call installed the wrapper.""" + global _original, _max_rows, _on_taken, _on_missed + device = _indexed(device) + if device not in _maps or _maps[device].shape[0] < max_rows: + t = torch.arange(max_rows, dtype=torch.int32, device=device) + _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() + _max_rows = max(_max_rows, max_rows) + _on_taken = on_taken + _on_missed = on_missed + if _original is not None: + return False + _original = thop.trtllm_gen_generation_preprocess + thop.trtllm_gen_generation_preprocess = _guarded + logger.info( + "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " + f"(up to {max_rows} rows)" + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107__tp2/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_107__tp2/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/__init__.py new file mode 100644 index 000000000000..e6c3d806a4a3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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 TP2 +projection shapes on SM107 (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..3f74be245ee0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # capacity: fixed: a kernel tile constant (weight rows per CTA = the MMA's M) +MMA_N = 8 # token columns +CTA_K = 128 # capacity: fixed: a kernel tile constant (one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row) +MMA_INST_K = 32 +# 4 MMAs per k-tile, one per scale byte of the atom. +K_BLOCKS = CTA_K // MMA_INST_K # capacity: fixed: a kernel tile constant +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], +) -> None: + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/op.py new file mode 100644 index 000000000000..3124827082df --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/op.py @@ -0,0 +1,496 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) at the TP2 per-rank projection shapes (SM107). + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring[, contig, nacc, a_tmem, st_async]) of the TP2 per-rank projections. The optional fields default +# to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. SM107 graph us at M 1 / 4 / 8 over weights rotated beyond L2, against the +# stock MXFP8 quantize + FlashInfer CuTe DSL GEMM: QKV 4.11 / 4.12 / 4.10 vs 8.6-8.8, o_proj 3.90 / 3.96 / 3.70 vs +# 6.9-7.2, shared gate_up 3.11 / 3.09 / 3.10 vs 8.1-8.4, shared down 2.22 vs 4.7-4.9, dense gate_up 8.15 / 8.56 / 8.55 +# vs 10.2-10.5, dense down 5.89 / 5.72 / 5.46 vs 8.7-9.0. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + # sparse-layer QKV + index projection (index values disabled); timed at N 5120: 4992 is 39 of its 40 N tiles + (4992, 6144): (4, 12), + # dense-layer QKV + (4608, 6144): (4, 12), + # o_proj + (6144, 4096): (2, 12), + # shared-expert gate_up + (3072, 6144): (4, 12), + # shared-expert down + (6144, 1536): (2, 12), + # dense-layer gate_up + (12288, 6144): (2, 12), + # dense-layer down + (6144, 6144): (2, 12), +} +# Per-token-count configs (N, K, M); none at TP2. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = {} +# Above 8 tokens (wider than the TP2 decode step); none at TP2. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = {} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +# (rule 1b, deterministic) The decode projections at up to STRIDED_MAX_ROWS rows split K over strided k-tiles per +# split rank (contig 0: rank r owns k-tiles r, r + split, ...) instead of contiguous ranges: another fixed fp32 order. +# SM107 TP2 4 rows (783556), us/call contiguous -> strided: sparse QKV 4.100 / 4.102 -> 4.074 / 4.072, dense QKV +# 3.820 / 3.816 -> 3.812 / 3.805, o_proj 3.970 / 4.017 -> 3.945 / 4.001, dense gate_up 8.694 / 8.665 -> 8.627 / 8.606 +# (dense down stays contiguous: strided is slower); deterministic (3 runs bitwise) and max error vs the fp64 product +# equal to the contiguous order's. Kill: False. +STRIDED_K_TILES = True +STRIDED_MAX_ROWS = 8 # capacity: fixed: the strided-input kernels' built rows +_STRIDED = frozenset({(4992, 6144), (4608, 6144), (6144, 4096), (12288, 6144)}) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + if ( + STRIDED_K_TILES + and m is not None + and m <= STRIDED_MAX_ROWS + and (n, k) in _STRIDED + and len(cfg) == 2 + ): + cfg = cfg + (0,) + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # 17..32 rows at K < 1024 decline (an unsplit K below 1024 is slower than the FlashInfer GEMM there); no TP2 + # projection has K < 1024. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + if nacc != 1: + # More than one TMEM accumulator is broken on SM107 (783556 / 783557 / 783579: nondeterministic or garbage + # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + ) + return fn + + +def _views( + x: torch.Tensor, xsf: Optional[torch.Tensor], weight, weight_scale, y, x_bf16: bool +) -> tuple: + k = x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + n = weight.shape[0] + if weight.shape[1] != k: + return None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar) + if resolved is None: + return None + key, runtime_ar = resolved + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros(m, k, dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, device=dev) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode.py new file mode 100644 index 000000000000..b383cd4036a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode.py @@ -0,0 +1,1747 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +GB300/SM103 tuning targets GQA group 16, top-k 16 and four-token Eagle3 +verification. Split-K, query pairing and launch tables require four KV heads; +the linear-softmax policy also covers one and two KV heads. Other hardware +and shapes retain the default policies, including supported SM100 NVFP4. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* A packed NVFP4 cache is read in place, dequantized in-register from its E2M1 + nibbles and E4M3 block scales. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* A single top-k chunk writes the final output directly, avoiding scratch and + the otherwise redundant LSE merge. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. + +This is the MiniMax-M3 NVFP4 / sm_107 / tp2 target's copy. On top of the +in-tree decode it routes long-context TP4 steps to the one-launch +``m3_attn_lc`` kernel, short-sequence verify steps of at least +``MQ_MIN_TOKENS`` rows to the multi-query producer, releases PDL dependents +early, and has the merge emit the MXFP8 o_proj hand-off (``m3_prequant.SLOT``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + check_decode_span_shape, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +from . import m3_attn_lc, m3_prequant, m3_report + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +_TL_DTYPES = { + torch.bfloat16: tl.bfloat16, + torch.float16: tl.float16, + torch.float32: tl.float32, +} + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size + +# NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale +# byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. +NVFP4_SF_VEC_SIZE = 16 + + +# Fewest query rows for which the merge writes the MXFP8 o_proj hand-off: every decode step, so the MXFP8 o_proj +# never runs its own quantize kernel after attention. +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# The sparse decode (LC2, the Triton merge) also writes the MXFP8 o_proj hand-off (bit for bit +# torch.ops.trtllm.mxfp8_quantize of its BF16 rows) and publishes it to m3_prequant.SLOT. +MERGE_MXFP8_HANDOFF = True +# The long-context decode (m3_attn_lc) stages the request's block-table row before its PDL grid wait: off (measured +# +1.6..+1.8 us at 32 rows; flat to -0.2 us below), up to LC_PREWAIT_MAX_ROWS query rows when on. +LC_PREWAIT = False +LC_PREWAIT_MAX_ROWS = 16 # capacity: fixed: measured occupancy crossover +# The long-context decode's shapes (modeling's SPARSE_DECODE_LC keys: "mxfp8" = a step of at most MergeMxfp8.MAX_ROWS +# rows) and the identity-selection producer's (modeling's SPARSE_DECODE_PRODUCER keys), with their report paths. +LC_KEYS = frozenset({"mxfp8"}) +LC_PATH = "m3_sparse_decode_lc" +PRODUCER_KEYS = frozenset({"per_row"}) +PRODUCER_PATH = "sparse_decode_producer" + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and + swizzled UE8M0 scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips + its own quantization. One buffer of ``MAX_ROWS`` rows (one 128-row scale tile) per width serves every layer: each + o_proj consumes it before the next layer's merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + @staticmethod + def handoff(mxfp8, rows: int): + """The hand-off of a ``rows``-row step: its E4M3 rows and the 128-row scale tiles that cover them. The o_proj + GEMM checks the swizzled scales against the step's padded rows, so the step gets the tiles that cover its rows, + never the whole buffer.""" + q8, sf = mxfp8 + width = q8.shape[-1] + return q8[:rows], sf[: (rows + 127) // 128 * 128 * (width // 32)] + + +MERGE_MXFP8 = MergeMxfp8() + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + if IDENTITY: + # Slot `it` of the identity selection, bounded by this row's own kv extent above. + blk = it + else: + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, + # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + if IDENTITY: + blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) + blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) + else: + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for cand in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + if IDENTITY: + blk_0 = tl.where(valid_0, idx_0, sentinel) + blk_1 = tl.where(valid_1, idx_1, sentinel) + else: + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the + # pool, reads nothing. + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. Because this trigger precedes this grid's own + # wait, it says nothing about what upstream has finished: a dependent may read only step-invariant data + # (weights) before its griddepcontrol.wait, never this step's attention output, q, top-k table or lengths. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) +# Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) + + +# Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query +# tokens; below it the per-token producer is as fast. +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover +# Warps per multi-query producer program. +MQ_NUM_WARPS = 8 + + +def _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, +) -> None: + """minimax_m3_sparse_attn_decode through m3_attn_lc: same table, output and MXFP8 hand-off contract.""" + use_pdl = get_env_enable_pdl() + total_q, num_heads, head_dim = q.shape + prewait = bool(LC_PREWAIT) and total_q <= LC_PREWAIT_MAX_ROWS + mxfp8 = None + m3_prequant.SLOT.clear() + if MERGE_MXFP8_HANDOFF and MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(num_heads * head_dim, output.device) + # The early trigger lets the o_proj launch and stream its weights while attention runs. + m3_attn_lc.m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + linear_softmax=False, + mxfp8=mxfp8, + use_pdl=use_pdl, + prewait=prewait, + n1_ids=m3_attn_lc.N1_IDS.take(), + ) + if prewait: + m3_report.taken(LC_PREWAIT_PATH, total_q) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), *MergeMxfp8.handoff(mxfp8, total_q) + ) + + +# Engagement path of the identity-selection producers reading their blocks from the kv lengths (IDENTITY). +DEVICE_IDENTITY_PATH = "sparse_decode_device_identity" +# Engagement path of LC2's pre-wait block-table staging (LC_PREWAIT, up to its row cap). +LC_PREWAIT_PATH = "sparse_decode_lc_prewait" + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + q may be FP8, as a fused QK-norm/RoPE producer emits it. The kernel widens + it to output.dtype in-register rather than making the caller materialize a + widened copy, which would cost a standalone kernel on every sparse layer. + The attention math is unchanged either way, since E4M3 widens exactly. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. + + Passing the four NVFP4 scale tensors selects the packed-NVFP4 cache instead: + k_paged/v_paged then hold two E2M1 elements per byte, the block scales are + paged E4M3 bytes in the layouts msa_scatter writes, and the two per-tensor + dequant scales are single-element FP32 device tensors. + + num_topk_chunks overrides the split-K factor and exists for tests: the + merged result must not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) + nvfp4_args = (k_block_scale, v_block_scale, k_global_scale, v_global_scale) + kv_nvfp4 = any(arg is not None for arg in nvfp4_args) + scale_cols = head_dim // NVFP4_SF_VEC_SIZE + if kv_nvfp4: + capability = torch.cuda.get_device_capability(q.device) + if capability not in _SM100F_CAPABILITIES: + raise NotImplementedError( + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; " + f"got SM{capability[0] * 10 + capability[1]}." + ) + _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) + # Two E2M1 elements per byte; Triton must see them unsigned so that the + # high-nibble shift does not sign-extend. + k_paged, v_paged = k_paged.view(torch.uint8), v_paged.view(torch.uint8) + k_block_scale, v_block_scale = ( + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused arguments. + scale_arg = kv_scale if use_scale else output + k_sc_arg, v_sc_arg = (k_block_scale, v_block_scale) if kv_nvfp4 else (output, output) + k_gs_arg, v_gs_arg = (k_global_scale, v_global_scale) if kv_nvfp4 else (output, output) + + widen_q = q.dtype in _FP8_DTYPES + if widen_q and output.dtype not in _TL_DTYPES: + raise ValueError( + f"MiniMax-M3 sparse decode cannot widen FP8 q into {output.dtype}; " + f"supported compute dtypes are {sorted(d.__str__() for d in _TL_DTYPES)}." + ) + # Packed NVFP4 dequantizes to FP16. The production producer supplies FP8 + # q, which widens exactly to FP16, so keep both dot operands there instead + # of converting every K/V tile to BF16. Accumulation remains FP32. + if widen_q: + compute_dtype = tl.float16 if kv_nvfp4 else _TL_DTYPES[output.dtype] + else: + compute_dtype = tl.float32 + + # Long-context TP2 steps: attention, split-K merge and MXFP8 hand-off in one launch. A selected table on an NVFP4 + # cache with FP8 q at the TP2 heads is the kernel's (modeling's SPARSE_DECODE_LC); operands it declines are a miss. + lc_key = "mxfp8" if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS else "plain" + lc_covered = ( + kv_nvfp4 + and widen_q + and not identity_blocks + and num_topk_chunks is None + and num_heads == 32 + and num_kv_heads == 2 + and lc_key in LC_KEYS + ) + if lc_covered and not m3_attn_lc.lc_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ): + m3_report.missed( + LC_PATH, + total_q, + "the step's q / KV pages / scales / top-k table / output are not the layout the kernel is built for", + ) + lc_covered = False + if lc_covered: + _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + ) + return + + query_group_size = ( + _sm100f_nvfp4_query_group_size( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else 1 + ) + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = ( + _sm100f_nvfp4_num_topk_chunks( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else None + ) + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one + # program per (request, block) covers all of them and dequantizes each block once. + use_mq = ( + identity_blocks + and total_q >= MQ_MIN_TOKENS + and kv_nvfp4 + and num_kv_heads == 1 + and decode_query_len > 1 + and head_dim == 128 + and gqa_group_size * decode_query_len <= 128 + ) + if identity_blocks and kv_nvfp4 and widen_q: + # The identity-selection route is modeling's SPARSE_DECODE_PRODUCER ("per_row"; "multi_query" needs one KV + # head, so it never runs at the TP2 heads); a variant the range does not certify is a miss (raises under + # require). + producer_key = "multi_query" if use_mq else "per_row" + if producer_key not in PRODUCER_KEYS: + m3_report.missed( + PRODUCER_PATH, + total_q, + f"the identity-selection producer variant {producer_key!r} is not certified at {num_kv_heads} KV heads", + ) + if use_mq: + num_topk_chunks = max_topk + direct_output = num_topk_chunks == 1 + if direct_output: + # With one chunk, acc_o is already normalized and final. Point the + # producer at the caller's output and compile away the LSE store. + # Keeping a real dummy pointer preserves Triton's pointer contract. + o_arg = output + lse_arg = output + o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) + lse_strides = (0, 0, 0) + else: + # Persistent arena rather than torch.empty, so the partials keep one + # address across CUDA graph replays. FP32 keeps the split-K factor from + # perturbing the merged result. + reserve = torch.cuda.is_current_stream_capturing() + o_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + o_strides = tuple(o_arg.stride()) + lse_strides = tuple(lse_arg.stride()) + + use_pdl = get_env_enable_pdl() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = ( + _sm100f_nvfp4_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + if total_q % query_group_size or decode_query_len % query_group_size: + raise ValueError( + f"query group {query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." + ) + + producer_kernel = ( + _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel + ) + device_identity = identity_blocks + if device_identity and not use_mq: + m3_report.taken(DEVICE_IDENTITY_PATH, total_q) + single_kernel_args = ( + { + # No split-K arrival counters: the separate merge kernel combines the chunks. + "merge": (output, None), + "tail": (output.stride(0), output.stride(1), output.stride(2), num_kv_heads), + # An early release pays off when the dependent is the consumer of the + # final output (it can prefetch), or a merge grid that itself releases + # its dependent at entry: the o_proj then launches during this kernel and + # streams its weights while attention runs. Above 8 query rows the + # concurrent weight stream slows the larger attention grid more than the + # prefetch saves (GB300 TP4: -1.1% at 48 rows). + "options": {"FUSED_MERGE": False, "EARLY_TRIGGER": direct_output or total_q <= 8}, + } + if query_group_size == 1 + else {"merge": (), "tail": (), "options": {}} + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + else {} + ) + if use_mq: + from .triton_sparse_decode_mq import launch_mq_producer + + launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + sm_scale=sm_scale, + decode_query_len=decode_query_len, + use_pdl=use_pdl, + early_trigger=False, + num_warps=MQ_NUM_WARPS, + ) + else: + producer_kernel[ + (total_q // query_group_size) * num_topk_chunks, + num_kv_heads, + ]( + q, + k_paged, + v_paged, + scale_arg, + k_sc_arg, + v_sc_arg, + k_gs_arg, + v_gs_arg, + topk_idx, + o_arg, + lse_arg, + *single_kernel_args["merge"], + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + k_sc_arg.stride(0), + k_sc_arg.stride(1), + v_sc_arg.stride(0), + v_sc_arg.stride(1), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + *o_strides, + *lse_strides, + block_table.stride(0), + block_table.shape[1], + k_paged.shape[0], + *single_kernel_args["tail"], + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + WIDEN_Q=widen_q, + COMPUTE_DTYPE=compute_dtype, + USE_PDL=use_pdl, + KV_NVFP4=kv_nvfp4, + USE_LINEAR_SOFTMAX=use_linear_softmax, + SCALE_COLS=scale_cols, + DIRECT_OUTPUT=direct_output, + # Identity selection from the lengths this kernel reads, not the staged table. + IDENTITY=device_identity, + **single_kernel_args["options"], + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + m3_prequant.SLOT.clear() + if not direct_output: + width = num_heads * head_dim + mxfp8 = ( + MERGE_MXFP8.buffers(width, output.device) + if MERGE_MXFP8_HANDOFF + and MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS + and head_dim == 128 + and output.dim() == 3 + and output.is_contiguous() + and output.dtype == torch.bfloat16 + else None + ) + mxfp8_args = ( + { + "q8_ptr": mxfp8[0], + "sf_ptr": mxfp8[1], + "num_sf_k_tiles": width // 128, + "MXFP8_OUT": True, + } + if mxfp8 is not None + else {} + ) + merge_launch_options = ( + _sm100f_nvfp4_merge_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_arg, + lse_arg, + output, + head_dim, + *o_strides, + *lse_strides, + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + # Release the o_proj at the merge's entry so it streams its weights during the merge. + EARLY_TRIGGER=True, + **mxfp8_args, + **pdl_launch, + **merge_launch_options, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, width), *MergeMxfp8.handoff(mxfp8, total_q) + ) + + +__all__ = [ + "MERGE_MXFP8", + "MERGE_MXFP8_MIN_TOKENS", + "MergeMxfp8", + "NVFP4_SF_VEC_SIZE", + "SPARSE_BLOCK_SIZE", + "_sm100f_nvfp4_merge_launch_options", + "_sm100f_nvfp4_num_topk_chunks", + "_sm100f_nvfp4_query_group_size", + "_sm100f_nvfp4_use_linear_softmax", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107__tp2/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_107__tp2/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py new file mode 100644 index 000000000000..46329cb9d50d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py @@ -0,0 +1,6464 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_107 / tp2 -- self-contained modeling code. + +MiniMax-M3's language model: 60 layers, hidden 6144, 64 query heads and 4 KV heads of dim 128. Layers 0-2 are dense +attention with a dense SwiGLU-OAI MLP; layers 3-59 are MSA block-sparse attention (an index branch scores 128-token +blocks and the top 16 are attended) with a MoE of 128 routed experts (top 4, sigmoid scores with a correction bias, +renormalized and scaled by 2) plus one shared expert. The checkpoint is a modelopt mixed-precision export: NVFP4 routed +experts, MXFP8 projections, fp32 router, bf16 embedding and head; the KV cache is NVFP4 for the sparse layers and FP8 +for the dense ones. + +``tp2`` is two Rubin GPUs: head-split attention (32 query heads, two KV heads and two index heads per rank, one +index-K head) and tensor-parallel experts (every expert on every rank, intermediate 1536 per rank), no attention DP. +Speculative decoding, when configured, is one-model Eagle3 (draft length 3): the BF16 drafter (``M3Eagle3Draft``) and +the BF16 target head, both at checkpoint precision; decode steps of up to 2 requests verify and draft through fused +kernels at that precision (``M3Eagle3Worker``). + +**The core** (``MiniMaxM3Tp2ModelingV2``) holds the weights (``weights.MODEL_WEIGHTS``) and the runtime resources the +steps use -- one MSA attention backend per layer, the TP all-reduce communicator, the RoPE table, the MXFP8 GEMM +dispatch state the engine tunes -- and dispatches each step, on the host, from its context count alone +(``_select_target``): + +* ``DecodeTarget`` takes every pure generation step, of 1..``WIDE_MAX_TOKENS`` rows (up to 48 requests, each one token + or a draft verify of 1 + 3 rows). Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids it runs the decode-step + kernels and picks each layer boundary's collective by the step's token bucket (``DECODE_PLANS``: tiny 1..8, b16 + 9..16, b32 17..32); it runs its own generic layer for the first layer when the fused embedding does not take the + step, and for every layer of a wider step, with the wide cells (the 48-row generic boundaries and MC finalize, the + tensor-core router, the Triton MoE finalize, the 33..48-row GEMVs, the chunked greedy verify). The fast-path cells + go up to 128 rows; a 129..192-row step takes the stock-exact paths. +* ``PrefillTarget`` takes every step with context rows: prefill chunks and mixed steps, every layer in its generic + form over the MSA backend. + +The two targets share no body: each owns its layer, attention and MoE compositions and its ``step_args``, the one +projection of the step metadata its forward reads. What both call are the module-level op adapters, each picking the +kernel for one call from shapes and layouts, which read no step metadata themselves. + +Every fast path the decode step can take, the rows each covers (``FAST_PATHS``, each paired with the entry whose +``CELLS`` certify it) and what a step that cannot take one does are below: **uncovered** (outside every covered +range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise +``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). A row range is empty unless a +landed cell has certified it at the TP2 geometry on Rubin. Every fast path a GEN rank takes prints a rank-tagged +engagement line once per (path, rows) (``taken``). Kernel modules report through ``kernels/m3_report``, which this +module wires to the same policy at import. + +**Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the +next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the +long-context decode for the o_proj. It publishes it in ``kernels/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +MXFP8 projection on a decode graph takes it, and every take clears the slot. The decode boundaries' push targets (the +o_proj / dense-down / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's exchange) +are passed explicitly from the layer to the producer and back. + +Compiled kernels (the MoE engine, the decode GEMVs with their push variants, the index scorer, the long-row select, +the full-context decode plans) are built before the executor's warmup; a forward never compiles one, and one it finds +missing is a miss. Other state a fast path needs (multicast workspaces, push descriptors) is built on the eager decode +steps; a CUDA-graph capture only reads it, so a capture that finds it missing is a miss (``capture_missed``). + +**What this target asserts rather than adapts**: SM 10.7, the topology and checkpoint (``routing.py``); the +MiniMax-M3 MSA sparse attention runtime with the fused five-way QKV/index projection, an FP8 indexer and an FP8 or +NVFP4 KV cache, no torch.compile and no piecewise prefill graphs (``_validate_runtime``); text-only input (the vision +tower is not built and a step carrying image input raises). +""" + +from __future__ import annotations + +import atexit +import copy +import dataclasses +import os +import threading +import time +from dataclasses import dataclass +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + resolved_mode, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, +) +from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + _ensure_draft_vocab_size, + _get_requested_draft_moe_backend, + _set_draft_kv_cache_quant_algo, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + TensorParallelMode, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import ( + eager_on_graph, + is_in_breakable_cuda_graph, +) +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen, get_model_extra_attrs +from tensorrt_llm._utils import mpi_rank, prefer_pinned +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import weights as _weights +from .kernels import m3_draft_step as m3_draft_step_kernel +from .kernels import ( + m3_drafter_attn, + m3_fold_tail, + m3_mc_allreduce, + m3_prequant, + m3_report, + m3_spec_tail, +) +from .kernels import m3_fullctx_decode as m3_fullctx_decode_kernel +from .kernels import m3_index_score as _index_score_kernel +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 import triton_sparse_decode as _sparse_decode_kernel +from .kernels.long_select import long_select_blocks +from .kernels.m3_moe_finalize import m3_moe_finalize +from .kernels.mx_gemv import op as mx_gemv +from .ops.activation.m3_silu_and_mul import m3_silu_and_mul +from .ops.attention.m3_draft_attention import m3_draft_attention +from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope +from .ops.attention.m3_fullctx_decode import m3_fullctx_decode +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_msa_attention import m3_msa_attention +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 +from .ops.attention.m3_qkv_insert import m3_qkv_insert +from .ops.attention.m3_select_blocks import m3_select_blocks +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_select_long_dsl import m3_select_long_dsl +from .ops.attention.m3_sparse_decode import m3_sparse_decode +from .ops.attention.m3_sparse_decode_lc import m3_sparse_decode_lc +from .ops.comm.m3_allreduce import m3_allreduce +from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm +from .ops.embedding.m3_embedding import m3_embedding +from .ops.gemm import m3_ugpu_head as m3_ugpu_head_op +from .ops.gemm.m3_bf16_linear import m3_bf16_linear +from .ops.gemm.m3_dense_swiglu import ALPHA as DENSE_ALPHA +from .ops.gemm.m3_dense_swiglu import BETA as DENSE_BETA +from .ops.gemm.m3_dense_swiglu import LIMIT as DENSE_LIMIT +from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu +from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS +from .ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for +from .ops.gemm.m3_fp32_linear import m3_fp32_linear +from .ops.gemm.m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + m3_mx_gemv, + served_max_rows, +) +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.gemm.m3_ugpu_head import m3_ugpu_head +from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .ops.spec.m3_capture_norm import m3_capture_norm +from .ops.spec.m3_chunk_norm import m3_chunk_norm +from .ops.spec.m3_draft_input import m3_draft_input +from .ops.spec.m3_draft_silu import m3_draft_silu +from .ops.spec.m3_draft_step import m3_draft_step +from .ops.spec.m3_force_vec import m3_force_vec, m3_force_vec_req +from .ops.spec.m3_sampled_accept import m3_sampled_accept +from .ops.spec.m3_shard_accept import m3_shard_accept +from .ops.spec.m3_spec_offsets import m3_spec_offsets +from .ops.spec.m3_step_prologue import m3_step_prologue + +_SM = (10, 7) +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The 1..8-row decode kernels' step bound: 2 requests verifying 1 + 3 drafts (the disaggregated GEN worker's batch). +TINY_MAX_TOKENS = 8 # capacity: fixed: the 1..8-row decode kernels' step bound +# Batch-4 plan: steps of 9..16 rows (3-4 requests verifying 1 + 3 drafts) also take the decode layer, on the "b16" +# plan below. Every fast path keeps its own rows (each up to TINY_MAX_TOKENS unless it names its 16-row rows), so a +# path not raised for 16 rows declines to the generic kernel inside the decode layer. False: every step above +# TINY_MAX_TOKENS rows runs the generic layer. +B4_DECODE = True +B4_MAX_TOKENS = 16 if B4_DECODE else TINY_MAX_TOKENS +# Batch-8 plan: steps of 17..32 rows (5-8 requests verifying 1 + 3 drafts) also take the decode layer, on the "b32" +# plan below; every path keeps the rows it names (a path raised only to 16 rows declines a 17..32-row step). False: +# the b16 plan's 16-row cap. +B8_DECODE = B4_DECODE +DECODE_MAX_TOKENS = 32 if B8_DECODE else B4_MAX_TOKENS +# The widest decode step this target takes (aggregated serving, up to 48 requests x (1 + 3 drafts) verify rows): decode +# rows 1..WIDE_MAX_TOKENS; the fast-path cells go up to 128 rows and a 129..192-row step takes the stock-exact paths +# (the Triton sparse-decode producer + merge, the generic GEMM). Steps past DECODE_MAX_TOKENS run the generic layer +# with the wide cells below (GENERIC_MC_*, ROUTER_TC_GEN, MOE_TRITON_FINALIZE, the W-G1 GEMVs, the wide verify). +WIDE_MAX_TOKENS = 192 # capacity: fixed: 48 requests x (1 + 3 drafts) + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 tp4 measurements in the comments).""" + + name: str + max_tokens: int + # The attention boundary's multicast all-reduce + residual + RMSNorm (5.05 vs 5.27 us native at 1 token). + attn_mc_ar: bool + # The same kernel at the MoE boundary with the shared + routed add folded in: wins up to 48 tokens. + moe_mc_ar: bool + # The MoE-boundary all-reduce also emits the next QKV activation in MXFP8 for the decode GEMV / GEMM. + qkv_prequant: bool + + +# tiny: every step up to 8 tokens (c1..c8 without speculation, the c1 verify, every draft pass); small: the c2 +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (kernels/m3_mc_allreduce.BUCKETS). +# The multicast all-reduce + residual + RMSNorm at both boundaries of every TP2 decode step (one 8-token bucket), the +# MoE boundary's also publishing the next QKV activation in MXFP8. SM107 TP2 hop: 3.5..4.4 us at 1..8 rows against +# 8.7..9.5 us for the stock AUTO (NCCL) all-reduce + fused residual RMSNorm, and 13.3..14.0 us at the MoE boundary +# (stock add + all-reduce + mxfp8_quantize); bitwise equal to stock (ops/comm/m3_mc_allreduce_norm). +# False puts both boundaries back on the built-in all-reduce. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +MC_AR_TP2 = True +# B3: the 9..32-row steps' attention and MoE boundaries on the same multicast kernel (the 16-token bucket at 9..16 rows, +# epilogue option 15; the 48-token bucket at 17..32: on TP2 the row kernel launches one CTA per row, so the bucket +# only sizes the exchange slots); bitwise equal to stock at every row count (the op's 9..32-row cells). False: the +# built-in all-reduce on those steps. +B4_MC_AR = MC_AR_TP2 +PLAN_MC_AR_MAX_TOKENS = DECODE_MAX_TOKENS if B4_MC_AR else TINY_MAX_TOKENS +# On those steps the MoE (past the engine's 8 rows) returns its unfinalized expert outputs and the multicast finalize +# kernel does the finalize + all-reduce + residual + RMSNorm in one launch (its cells cover 9..48 rows bitwise) +# instead of the routed finalize then the boundary. False: the routed finalize, then the boundary. +B4_MC_FINALIZE = B4_MC_AR +MC_FINALIZE_MIN_TOKENS = 9 # capacity: fixed: past the MoE engine's 8 rows +# B7 (SpecTail TR-5b): on the 9..32-row steps the boundary all-reduce (and the MoE's multicast finalize) also publishes +# the next QKV's MXFP8 activation (its quant cells at 9..32 rows are bitwise mxfp8_quantize), so QKV takes the MXFP8 +# GEMV on the producer-quantized x (795376: QKV at 16 rows 4.78 vs 2.59 us/call at 8; b8 census 794573: the generic +# MXFP8 GEMMs 1,196 + quantize 256 us/step at c8 vs 537 for the GEMVs at c2). +B4_QKV_PREQUANT = B4_MC_AR +B8_QKV_PREQUANT = B4_MC_AR +DECODE_PLANS = ( + ( + DecodePlan( + "tiny", + TINY_MAX_TOKENS, + attn_mc_ar=MC_AR_TP2, + moe_mc_ar=MC_AR_TP2, + qkv_prequant=MC_AR_TP2, + ), + ) + # b16: the 9..16-row steps; their boundaries on the multicast all-reduce's 16-token bucket. + + ( + ( + DecodePlan( + "b16", + B4_MAX_TOKENS, + attn_mc_ar=B4_MC_AR, + moe_mc_ar=B4_MC_AR, + qkv_prequant=B4_QKV_PREQUANT, + ), + ) + if B4_MAX_TOKENS > TINY_MAX_TOKENS + else () + ) + # b32: the 17..32-row steps; their boundaries on the 48-token bucket. + + ( + ( + DecodePlan( + "b32", + DECODE_MAX_TOKENS, + attn_mc_ar=B4_MC_AR, + moe_mc_ar=B4_MC_AR, + qkv_prequant=B8_QKV_PREQUANT, + ), + ) + if DECODE_MAX_TOKENS > B4_MAX_TOKENS + else () + ) +) + +# The dense layers' decode attention (trtllm-gen, msa_decode.dense_decode_max_seq_len): the max_seq_len it is handed +# per request count. It only sizes the multi-CTA KV split (every CTA still covers its share of the actual length; the +# result changes only by the split-merge order), and the model maximum over-splits AgentX-length caches at one +# request: GB300 TP4, 16 heads, FP8 KV, 64K context: 15.1 us at 16K vs 28.6 us at 1M. Two requests are neutral (the +# grid is unchanged) and keep the step's bound. + + +# TP2 skeleton: no cap (the built-in split) until a TP2 cap is measured and certified. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) + +# The dense layers' decode attention on PrimTS split-KV decode (kernels/m3_fullctx_decode) instead of the trtllm-gen +# call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the +# standalone reducer. Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 +# (29.3), (1, 1) 15.7 (28.0), (4, 2) 24.5 (29.3), (1, 2) 20.2 (28.1). The stock split is capped at 16 cluster CTAs +# per (request, KV head), i.e. 32 of 212 CTAs at one request. At one request (the E2E batch) the K/V stream is +# latency-bound, so the b1 recipes keep 12 K/V stages in flight (the load pipeline only; the same bits): locked +# clock, cold L2, 165k, (4, 1) 22.6 vs 26.9 and (1, 1) 20.7 vs 23.0 µs with the default 4 stages. +FULLCTX_DECODE_RECIPES: Mapping[Tuple[int, int], Mapping[str, Any]] = MappingProxyType( + { + (4, 1): MappingProxyType( + { + "config": {"tile_size_q": 32, "groups_tokens_heads_q": True, "kv_stages": 12}, + "splits": 48, + } + ), + (1, 1): MappingProxyType({"config": {"kv_stages": 12}, "splits": 72}), + (4, 2): MappingProxyType( + {"config": {"tile_size_q": 32, "groups_tokens_heads_q": True}, "splits": 24} + ), + (1, 2): MappingProxyType({"config": {}, "splits": 40}), + } +) +# False puts the dense layers back on the trtllm-gen call at every shape. +FULLCTX_DECODE_ENABLED = True +# The dense layers' FP8 q/k/v (decode steps of the hybrid cache): the TP2 op rounds Q/K to BF16 and then to E4M3 in +# its store, the bits of the BF16 fused_qk_norm_rope output cast to FP8. False runs that BF16 kernel and a separate cast +# (the same bits, one more kernel). +DENSE_QKV_FP8_VIA_BF16 = True +# On a full-context decode step that op also stores the new K/V rows into the layer's paged pool (the bytes +# write_msa_phase_kv's indexed copies store, one launch instead of about ten). False writes them with +# write_msa_phase_kv. +DENSE_KV_INSERT_FUSED = True +# The engagement line of that fused store (dispatch.taken, one per row count and rank). +DENSE_KV_INSERT_PATH = "m3_dense_kv_insert" +# The drafter MLP's C++ SwiGLU (DRAFT_SILU_CPP, spec_worker._m3_draft_mlp): its own taken / not-taken lines. +DRAFT_SILU_PATH = "m3_draft_silu" +# The split-KV merge on the BF16 op (minimax_m3_sm107_tp2_fullctx_reduce) instead of PrimTS's FP16 reducer and a cast: +# Rubin, 165k, graph-us per call (1, 1) 14.5 vs 18.5, (4, 1) 17.3 vs 21.3, (1, 2) 18.5 vs 23.1, (4, 2) 23.3 vs 27.5. +FULLCTX_DECODE_OWN_REDUCER = True +# The K/V length bound the full-context plans are compiled for; a step whose cache length bound is longer is a miss. It +# covers the TP2 AgentX config's max_seq_len (1 << 20) plus the speculative rows the cache reserves past it. +FULLCTX_DECODE_MAX_KV_LEN = 1 << 21 # capacity: fixed: the full-context plans' compiled K/V bound +# The Eagle3 drafter layer's generation phase on the same full-context decode (kernels/m3_drafter_attn: the in-tree +# RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its +# split at 8 CTAs per (request, KV head) on SM 10.7: 45-46 us per draft pass at 165k. +DRAFTER_FULLCTX_ENABLED = True +# On those passes the drafter's RoPE, FP8 query and FP8 K/V append run as one launch (the op +# minimax_m3_sm107tp2_drafter_rope_fp8_kv, a copy of the certified TP4/TP8 kernel at the TP2 heads: the in-tree +# preprocess's GPT-NeoX rotation and unit-scale E4M3 store, and the torch E4M3 cast of its BF16 query) instead of +# thop.trtllm_gen_generation_preprocess and the cast (3.6 vs 1.6 us per pass, 792311). +DRAFTER_ROPE_FUSED = True +# The engagement line of that launch (dispatch.taken, one per row count and rank). +DRAFTER_ROPE_PATH = "m3_drafter_rope" +# How the rotation's products are summed, as the in-tree preprocess build sums them (0: two products then the add, +# 1: fma(cos, x, sin * pair), 2: fma(sin, pair, cos * x)); the op's bitwise test pins it. +DRAFTER_ROPE_MODE = 1 +# thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request +# multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by +# that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, +# token-0 drafts). kernels/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# calls: what the build writes (padded rows included), so bitwise the op with a fresh map. False:. +PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the index scorer launches and streams its pre-wait index-K prefetch while the producer +# runs (its outputs are read only after the scorer's own grid wait). Rubin TP2, 57-layer PDL chain behind a 21 MB +# pre-wait L2 prefetch (the 165K scorer's shape): 3.59 (native) -> 1.93 us per layer at 4 rows, 3.64 -> 2.00 at 8. +PRODUCER_EARLY_TRIGGER = True +# The producer also reads its per-step inputs (position ids, cache slots, RoPE rows) before its grid wait, leaving only +# the packed QKV row after it. Their writers are host-to-device copies before the replay and plain-launched kernels +# earlier in the step (the input preprocessing, msa_kv_update or its torch fallback), so they are final before the +# producer can start; a PDL-launched writer of any of them would have to turn this off. The call site's +# position_ids.reshape(-1).contiguous().to(torch.int32) and packed.contiguous() launch a plain torch copy right before +# the producer when the dtype or layout differs: also a plain writer, so the contract holds. Bit-identical outputs. +# PRE-WAIT WRITER CONTRACT (rubin-20261015/PREWAIT_AUDIT.md item b): the kernels that read step inputs (seq_lens, +# block tables, positions, cache slots, the sampled-accept pre-wait parameters) before their grid-dependency wait rely +# on every writer of them being a host-to-device copy before the replay or a PLAIN launch earlier in the step: +# msa_kv_update, the step-input prologue (Triton scatter or the C++ op), the spec fold tail (Triton or C++), the pack +# scatter and any step-head op (the spec step correction, STEP_SPEC_OFFSETS_ONE_LAUNCH). Never give one of them a PDL +# attribute with an early trigger. +PRODUCER_PREWAIT_STEP_INPUTS = True + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. SM107: 1.55..1.65 us against +# 3.4..6.0 us for the stock embedding + RMSNorm at 1..8 rows (the gather bitwise, the norm one ulp off on ~5e-6 of the +# elements: summation order); it also hands the first layer its residual, +# so layer 0 takes the decode layer instead of the generic one. False runs layer 0 on the generic layer. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +EMBED_FIRSTNORM_TP2 = True +# B5 (SpecTail TR-5b): the same kernel on the 9..16-row steps (one CTA per row, so each row's arithmetic is the 1..8-row +# one; the op's 9..16-row cells), which hands layer 0 its residual so it takes the decode layer. Not past 16 rows: +# there the stock embedding + RMSNorm it replaces differs from it by more than at 1..8 rows (7..21 normed elements, one +# past 1 ulp, at 19..28 rows; gate 795453 / debug 795562), so 17..32-row steps keep the stock layer 0. +# False: 1..8 rows. +B4_EMBED_FIRSTNORM = True +EMBED_FIRSTNORM_MAX_TOKENS = ( + (B4_MAX_TOKENS if B4_EMBED_FIRSTNORM else TINY_MAX_TOKENS) if EMBED_FIRSTNORM_TP2 else 0 +) +# The o_proj pushes its rows into the attention boundary's multicast all-reduce (prepushed mode 2) up to this many +# tokens; wider steps run the GEMV and a plain multicast all-reduce. SM107 TP2: -0.5..-1.2 us per layer, bitwise equal +# to the unpushed boundary. False keeps the GEMV and the boundary separate. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +ATTN_PUSH_TP2 = True +# B6 (SpecTail TR-5b): the o_proj also pushes on the 9..16-row steps, into the boundary's 16-token bucket; the GEMV's +# push epilogue covers MMA width 16 with the bucket's tokens in its compile key, and the pushed boundary is bitwise the +# unpushed one (the op's prepushed 9..16-row cells). Not past 16 rows: the o_proj GEMV keeps 16 rows +# (_GEMV_B4_ROWS_ONLY), so a 17..32-row o_proj runs the generic GEMM. +# False: 1..8 rows. +B4_ATTN_PUSH = B4_MC_AR +ATTN_PUSH_MAX_TOKENS = ( + (min(PLAN_MC_AR_MAX_TOKENS, B4_MAX_TOKENS) if B4_ATTN_PUSH else TINY_MAX_TOKENS) + if (MC_AR_TP2 and ATTN_PUSH_TP2) + else 0 +) +# The dense layers' MLP down projection pushes its rows into the post-MLP boundary's multicast all-reduce (prepushed +# mode 2), as the o_proj does into the attention boundary: the all-reduce no longer waits for the whole down GEMV grid +# before loading and pushing the rows itself. SM107 TP2 (stack-ON nsys 773230): the two non-capture dense layers' +# post-MLP hops end 5.5-5.8 us after the down GEMV vs 2.4-2.5 us for a pushed hop. Bitwise equal to the unpushed +# boundary (the same bf16 words). False keeps the down GEMV and the boundary separate. +DENSE_PUSH_TP2 = True +DENSE_PUSH_MAX_TOKENS = TINY_MAX_TOKENS if (MC_AR_TP2 and DENSE_PUSH_TP2) else 0 +# The generic layer's boundaries (prefill, mixed steps, Eagle3 capture layers, decode steps past DECODE_MAX_TOKENS): the +# most rows that take the multicast all-reduce + residual + RMSNorm, alone and with the shared + routed add folded in: +# the multicast all-reduce's widest bucket (48 tokens), as the agg canon runs it (v1.8.3; GSM T0 / T1 803743 / +# 803744), under the same constant (MC_AR_TP2; off = the built-in all-reduce). +GENERIC_MC_AR_MAX_TOKENS = 48 if MC_AR_TP2 else 0 # capacity: fixed: the widest bucket +GENERIC_MC_AR_FOLDED_MAX_TOKENS = 48 if MC_AR_TP2 else 0 # capacity: fixed: the widest bucket +# The generic layer's MoE boundary past DECODE_MAX_TOKENS rows, up to this many (at most the finalize kernel's 48, +# m3_mc_allreduce.FINALIZE_MAX_TOKENS): the routed experts stay unfinalized and their finalize (+ the shared expert) +# runs inside the boundary's multicast all-reduce + residual + next RMSNorm (m3_mc_finalize_allreduce_norm), as the +# decode layer's MC_FINALIZE does. 48 as the agg canon runs it (v1.8.3; the same gates as above). +GENERIC_MC_FINALIZE_MAX_TOKENS = 48 if MC_AR_TP2 else 0 # capacity: fixed: the finalize rows +# The MoE finalize (+ the shared expert add) of a finalized MoE forward with at least MOE_TRITON_FINALIZE_MIN_TOKENS +# rows (prefill, mixed and wide decode past the folds) runs as kernels/m3_moe_finalize: the trtllm-gen runner returns +# its FC2 output unfinalized and one Triton kernel gathers the top-k rows together, bitwise to the runner's own +# finalize and the separate shared add (153 -> 98 us per layer at 16K rows on SM107). False: the runner finalizes. +MOE_TRITON_FINALIZE = True +MOE_TRITON_FINALIZE_MIN_TOKENS = 49 # capacity: fixed: past the generic MC finalize's 48 rows +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch, ops/moe/m3_moe_engine at +# the TP2 width 1536). Rubin TP2 ablation chain (57 layers behind a 58.4 MB PDL producer, row = (P + MoE) - P, both +# orders): 31.3 / 50.7 us/layer at 4 / 8 tokens against the built-in MoE's 42.4 / 55.4 (FlashInfer MXFP8 shared +# expert); bit for bit the built-in block's shared.add_(routed). KILL_MOE_ENGINE: the built-in MoE at every row count. +KILL_MOE_ENGINE = False +MOE_ENGINE_MAX_TOKENS = 0 if KILL_MOE_ENGINE else 8 +# The engine's push of its sum into the MoE boundary's multicast all-reduce (prepushed mode 2) on the rows the +# engine serves, with the multicast boundary (DECODE_PLANS' moe_mc_ar): the boundary reads the pushed words instead +# of loading and pushing the rows itself. Bit for bit the unpushed boundary (comm/_m3_moe_push_op_matrix.py). +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS if MC_AR_TP2 else 0 +# When the engine pushes its sum into the boundary all-reduce, every CTA releases the PDL dependents right after its +# grid dependency wait instead of after the push, so the all-reduce (which polls the pushed tagged words, not the grid) +# becomes resident on the SMs the engine frees. Scheduling only: the same bits (the push oracle checks on vs off). +# Rubin TP2 greedy early column, same-job A,B + B,A (778191 / 778192): c1 -15 us, about -15 at 323K (within noise). +# False: releases after the push. +MOE_EARLY_AR_RELEASE = True +# A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel +# over the target lm_head's BF16 shard logits (kernels/m3_shard_accept), at the configured draft length. +SHARD_ACCEPT_DRAFT_LEN = 3 # capacity: fixed: the built Eagle3 draft length +# B4 spec tail: the fused verify (shard / sampled accept), the fused first drafter inputs (chunk / capture norm) and the +# fused draft loop (draft_step, its draft input, the drafter boundary all-reduce) also take the 9..32-row steps, up to +# DECODE_MAX_TOKENS verify rows (DECODE_MAX_TOKENS // (draft length + 1) requests; the draft step's head takes at most +# 8). Every kernel is row-local (per-row norms, per-request keys and Philox draws), so each row gets the bits it gets +# at 1 / 2 requests; the op cells at the new rows certify it. +# False: TINY_MAX_TOKENS rows. +B4_SPEC_TAIL = True +SPEC_TAIL_MAX_TOKENS = DECODE_MAX_TOKENS if B4_SPEC_TAIL else TINY_MAX_TOKENS +SPEC_TAIL_MAX_REQUESTS = SPEC_TAIL_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) +# W2 (SpecTail, agg wide decode): the greedy fused verify (kernels/m3_shard_accept) also takes the +# 33..WIDE_MAX_TOKENS-row decode-only steps, as whole-request chunks of the kernel's certified 48-row geometry (12 +# requests per launch at draft length 3; m3_shard_accept.WIDE_SPEC_TAIL), instead of the gathered target logits + +# argmax; and the in-tree draft loop takes the drafter's top-1 from its local head shard (W2b). Sampled verifies keep +# the fused path only up to SAMPLED_VERIFY_MAX_REQUESTS (past it, the in-tree sampler). Every other fused piece keeps +# SPEC_TAIL_MAX_TOKENS rows (the drafter then runs its in-tree loop). Rubin TP2 W1b 797141 (verify + 3 drafter top-1 +# passes, stock -> fused). Greedy chunks and the shard top-1 are bitwise the gathered argmax (lowest index on ties). +WIDE_SPEC_TAIL = m3_shard_accept_kernel.WIDE_SPEC_TAIL +WIDE_VERIFY_MAX_TOKENS = WIDE_MAX_TOKENS if WIDE_SPEC_TAIL else SPEC_TAIL_MAX_TOKENS +WIDE_VERIFY_MAX_REQUESTS = WIDE_VERIFY_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) +# 0: (the gathered logits + argmax run). Rubin TP2 greedy verify (microbench): stock 120.4 us at b1; fused +# with the coded head 117.0 us; the BF16-head + shard-accept arm is not measured yet. +SHARD_ACCEPT_MAX_REQUESTS = WIDE_VERIFY_MAX_REQUESTS +# A sampled (temperature + top-p, advanced_sampling_mode no_topk) Eagle3 verify step of up to this many requests +# verifies with the fused sampled acceptance (kernels/m3_sampled_accept, the TP2-shard kernel) over the target +# lm_head's BF16 shard logits, at the configured draft length (the AgentX contract). 0: (the gathered logits +# and the in-tree sampler run). Rubin TP2 b1/b2 (p5 microbench): 180.1 / 181.8 -> 115.1 / 113.9 us per verify. +SAMPLED_VERIFY_MAX_REQUESTS = SPEC_TAIL_MAX_REQUESTS +# A step verified by the in-tree path takes the fused first drafter inputs (the capture norms and the first pass's +# gather) up to this many requests. +FIRST_DRAFTER_INPUTS_MAX_REQUESTS = 0 # capacity: fixed: off at TP2 (0) +# A step verified by the fused kernels (dispatch.SHARD_ACCEPT / SAMPLED_ACCEPT) of up to this many requests takes the +# fused first drafter inputs (spec_worker._m3_first_drafter_inputs: the accepted tokens as they lie, the first pass's +# gather rows, the fc_norm chunk norm or the armed capture norm, then the in-tree fc). 0: (the in-tree first +# drafter inputs run, the captures stay in-tree and the fused draft loop steps aside). +FUSED_FIRST_INPUTS_MAX_REQUESTS = SPEC_TAIL_MAX_REQUESTS +# On a fused-verify step of up to this many verify rows the Eagle3 captures write rmsnorm(bf16(h + residual)) * +# fc_norm[i] in one kernel (ops/spec/m3_capture_norm) once the fc_norm weights are stacked; the fused first drafter +# inputs then skip their chunk norm. 0: (the in-tree capture, then the chunk norm). +CAPTURE_NORM_MAX_ROWS = SPEC_TAIL_MAX_TOKENS +# The first draft pass's per-layer fc_norm over the captured columns in one kernel (ops/spec/m3_chunk_norm) at up +# to this many verify rows; wider first passes take the in-tree first drafter inputs. 0:. +CHUNK_NORM_MAX_ROWS = SPEC_TAIL_MAX_TOKENS +# The fused draft loop's two all-reduce boundaries per pass (o_proj -> post-attention norm, down -> final norm) of up +# to this many rows run the multicast all-reduce + residual + RMSNorm (kernels/m3_mc_allreduce, opts bit 16: the +# drafter's stock rounding, normalizing the unrounded FP32 sum as FlashInfer fused_add_rmsnorm does) instead of the +# built-in all-reduce + RMSNorm. 0:. Rubin TP2 (p16 microbench): 10.0..11.0 -> 4.6..5.3 us per boundary. +# (The drafter epilogue, opts bit 16, is built for the 8-token bucket only: a B4 spec-tail first pass of 16 / 32 rows +# takes the built-in all-reduce + RMSNorm; its later passes of up to 8 rows stay multicast.) +DRAFT_MC_AR_MAX_ROWS = 8 # capacity: fixed: the drafter epilogue's 8-token bucket +# The drafter's o_proj and down projection, on the copied split-K kernel (DRAFT_LOWM_OWN_SPLITK), also push their rows +# into the drafter boundary all-reduce's exchange (the canonical words its own push would store, at the call parity), +# which then runs prepushed 2 instead of pushing x itself; the copy releases its dependents only after its own grid +# wait and the pushes (the prepushed >= 2 contract). Bitwise. TP2: the 4-row first pass only (its 1-row linears run +# FlashInfer's direct kernel). False:. +DRAFT_AR_PUSH = True +# The drafter boundary all-reduce fed by its projection's push (DRAFT_AR_PUSH): its own taken / not-taken lines. +DRAFT_AR_PUSH_PATH = "m3_draft_ar_push" +# The same push from the drafter's direct ('d') low-M linears (the 1-row passes' o_proj and down at TP2), through this +# target's copy of FlashInfer's direct kernel (kernels/m3_lowm_direct; its output FlashInfer's): the owning thread +# stores the CTA's column pairs into the exchange before the kernel's last-instruction release. Bitwise. Rubin TP2 +# world 2 (dpush-1 782685), per 1-row boundary: o_proj 9.22 -> 8.09 us, down 13.81 -> 12.72. False:. +DRAFT_AR_PUSH_DIRECT = True +# uGPU: a verify of these row counts computes the target logits shard as two half-vocabulary in-tree GEMMs, each on +# the green context of the die whose memory holds that half of a placed copy of the head (ops/gemm/m3_ugpu_head; the +# module's weight is untouched, so every other reader keeps ordinary memory); the logits are the in-tree head's bit for +# bit. Rubin TP2 head shard: 93.3 -> 79.5 us at +# 4 rows, 93.5 -> 80.6 / 80.7 at 12 / 16 (at 1 / 2 rows the half-shard GEMM is another kernel). In-model the decode +# step is ~45 us shorter at one request and ~90 us at two, and the gain holds at long context. +KILL_UGPU_HEAD = False +UGPU_HEAD_ROWS: Tuple[int, ...] = () if KILL_UGPU_HEAD else (4, 8, 12, 16) +# On a fused-verify step of up to this many verify rows the drafter's first fc is forked onto the aux stream right +# after the last captured target layer, overlapping the remaining target layers and the head (spec_worker). +FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover +# A sampled fused verify of up to this many verify rows on the in-tree head GEMM forks the drafter's first fc after the +# head instead of beside it (spec_worker): the fc's split-K grid otherwise holds the SMs first and the head GEMM starts +# late. 0:. Rubin TP2 c1 (locked nsys 768559): the head started 17.4 us after its input was ready. +FC_AFTER_HEAD_MAX_ROWS = SPEC_TAIL_MAX_TOKENS +# A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (spec_worker): per pass the +# draft input in one launch (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_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +# 0: (the in-tree draft loop runs). Rubin TP2 per draft pass (head + argmax + exchange): 122.6 / 126.0 -> +# 102.8 / 102.8 us at b1 / b2 (BF16 fragments; the coded head is slower on sm_107: DRAFT_HEAD_CODED False). +DRAFT_STEP_MAX_REQUESTS = min(SPEC_TAIL_MAX_REQUESTS, 8) # capacity: fixed: the draft head's 8 rows +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) run FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm, the TP8 port) at the rows its TP2 table is certified for, up to this many rows; other calls +# keep the in-tree cuBLAS GEMM. 0:. Rubin TP2 (mclk locked, 4 rotating sets) at 1 row: qkv 13.18 -> 9.71, +# o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. +DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the certified low-M rows +# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (kernels/m3_lowm_splitk): bit for bit the +# same output at every certified cell (spec/test_modeling_v2_m3_lowm_splitk.py); the seam for the drafter-boundary push +# epilogue and the down projection's SwiGLU prologue. False: (FlashInfer's kernel). +DRAFT_LOWM_OWN_SPLITK = True +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(m for m in rows if m <= DRAFT_LOWM_MAX_ROWS) + for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() +} +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (kernels/triton_sparse_decode.MQ_MIN_TOKENS); the multi-query producer needs one KV head, so at the TP2 heads every +# identity step takes the per-row producer. +SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover +SPARSE_DECODE_KV_HEADS = 2 # capacity: fixed: the tp2 KV heads per rank +# The long-row top-16 selector takes engines of 129..LONG_SELECT_MAX_BLOCKS KV blocks of 128 tokens; it holds 2, 4 or +# 9 keys per thread by the row's valid-block count. At most 128 blocks, or past the maximum, the built-in selector +# runs; 0 turns the path off. Rubin TP2 (2 KV heads), 165,521 tokens, behind the scorer, ablation per layer: +# -0.9 us at 1 request and -1.2 us at 2 requests against the built-in selector. +LONG_SELECT_MIN_BLOCKS = 129 # capacity: fixed: past the built-in selector's 128 +LONG_SELECT_MAX_BLOCKS = 9216 # capacity: fixed: the select kernels' built blocks per row +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +# The long-row select releases the sparse decode after it right after its own grid wait (the decode stages q and the +# request length before its wait) up to this many query rows; past it the trigger stays at the select's exit. On Rubin +# TP2 the early-launched long-context sparse decode (SPARSE_DECODE_LC) crowds the select's CTAs off the SMs already at +# one request (4 rows): releasing at the exit is 336 us per decode step faster at 1 request and 350 at 2 (in-model, +# graph time, 165K-context mixed draws, locked memory clock), so the early release is off at every row count. +SELECT_EARLY_MAX_ROWS = 0 # capacity: fixed: measured early-release crossover +# Without the early release, the long-row select releases the sparse decode once every score of its row is loaded +# (1, before the exact rank of the candidates) instead of at its exit (0). The release still follows +# the select's own grid wait, which the sparse decode's pre-wait q / length / block-table reads rely on. The decode's +# 16-CTA cluster launch and prologue then overlap the rank: Rubin TP2 in-model, locked memory clock, AgentX draws, +# 2 instances vs the r4.3 stored A: replay period -162 us fixed, -94 at a 323K-token request. +SELECT_LONG_TRIGGER_PHASE = 1 +# The long-row select runs as kernels/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per +# (query row, KV head) instead of one 1024-thread CTA, the same selection bit for bit (ties included), with 4, 11 or +# 24 keys per thread (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices. Rubin TP2 decode chain (2 KV heads, +# 4 rows, 57 layers behind the producer stream and the scorer, both orders x 3 rounds): -0.127 / -0.148 / -0.037 us +# per layer at 98K / 165K / 323K tokens against the 1024-thread kernel. False: (the 1024-thread kernel). +SELECT_LONG_DSL = True +LONG_SELECT_DSL_KEY_WIDTHS = (4, 11, 24) +# The MSA index scorer (ops/attention/m3_index_score) of a sparse decode step, up to INDEX_SCORE_MAX_REQUESTS requests +# (0 turns the path off; wider steps run the in-tree CuTe DSL scorer): its CTAs per request by request count, and up +# to INDEX_SCORE_EARLY_MAX_REQUESTS requests the early issue (Q and the first K tile right after the grid wait, one +# older block per CTA prefetched into L2 before it). The split is captured with the graph, so it is one value per +# request count over every context. Rubin TP2 (2 index heads), 165,521 tokens, behind a 58 MB PDL producer, ablation +# per layer: 1 request 3.27 us at 1536 CTAs (512: 3.76, in-tree 256-CTA scorer 6.07); 2 requests 4.97 us at 768 +# (in-tree 6.71). Within 0.2 us of each context's best split from 16K to 400K tokens except 2 requests at 400K (+0.6). +# The early issue reads the request's first page id from the MSA page table before the grid wait: exact because the +# page table is written only on the host side before the replay (attention/test_modeling_v2_m3_block_table_contract.py); +# the request length read there is a hint, every load and store is bounded by the one read after the wait. +# The batch-8 plan's 3- and 4-request steps take it too (Indexer k11 795543: Rubin TP2, 512 MB flush before each call, +# graph-timed, bitwise equal to the in-tree scorer; requests at 1.0 / 0.9 / 0.8 / 0.7 x the context): 3 requests at +# 424 CTAs each 6.08 / 9.04 / 14.20 us at 98K / 165K / 323K tokens (in-tree 6.65 / 9.44 / 14.49); 4 requests at 318 +# 7.41 / 10.82 / 16.80 (in-tree 7.60 / 10.97 / 16.96). 5..8 requests at about 1272 CTAs in total, 1272 // requests +# each (Indexer k11c 795819): 5 at 254 8.25 / 11.44 / 19.10 us (in-tree 8.66 / 12.10 / 19.65); 6 at 212; 7 at 181; +# 8 at 159 9.91 / 13.97 / 23.33 (10.33 / 14.47 / 23.74). Each within 0.03 us of its count's best split. +# 9..48 requests (W-IDX1, one split per bucket) win standalone (Indexer k14b 797357: 46 requests 52.29 -> 49.85 us per +# layer) but lose in-model (agg census c46 / c20: 52.13 -> 61.85 / 24.22 -> 29.06 us per layer, the in-tree scorer +# already at the single-stream bandwidth floor), so steps past 8 requests keep the in-tree scorer. +INDEX_SCORE_MAX_REQUESTS = 8 # capacity: fixed: measured in-model crossover +#: (largest request count of the bucket, CTAs per request) for 9..48 requests (W-IDX1; outside INDEX_SCORE). +INDEX_SCORE_WIDE_SPLITS = ((10, 127), (13, 97), (17, 74), (23, 55), (31, 41), (39, 32), (48, 26)) +INDEX_SCORE_SPLIT_K: Mapping[int, int] = MappingProxyType( + {1: 1536, 2: 768, 3: 424, 4: 318, 5: 254, 6: 212, 7: 181, 8: 159} + | {r: next(s for hi, s in INDEX_SCORE_WIDE_SPLITS if r <= hi) for r in range(9, 49)} +) +INDEX_SCORE_EARLY_MAX_REQUESTS = 48 # capacity: fixed: measured early-issue crossover + + +# Synth only (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS, fractional): the fused verify's forced accepted counts are +# drawn per request, keyed by the request's KV length at the verify (its prompt plus what it has accepted), instead of +# from the server-wide counter, so a request's forced counts do not depend on what the server ran before (A/B draws +# become history-independent). Production sampling never forces counts. False: (the shared counter, the in-tree +# sequence). +FORCE_AL_PER_REQUEST = True + +# The decode MXFP8 GEMVs (kernels/mx_gemv, SM107 TP2 table), per projection (N, K): the rows each serves, on an +# activation its producer quantized (the MoE boundary for the QKV, the attention merge for the o_proj) and on a bf16 +# one (quantized in the GEMV's prologue, bit-identical to mxfp8_quantize). SM107 graph us per projection at 1..8 rows: +# 2.2..8.6 against 4.7..10.5 for the stock MXFP8 quantize + FlashInfer CuTe DSL GEMM (kernels/mx_gemv/op.py). +# False puts every projection back on the stock decode-graph GEMM. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +MX_GEMV_TP2 = True +# MX-TAC: the decode-graph MXFP8 projections the GEMVs do not serve take one fixed per-bucket decision instead of the +# AutoTuner's per-start choice (whose L2-hot profiling left c2 bimodal at +-6%): FlashInfer's CuTe DSL quantize and +# CuTe DSL GEMM. Stock's candidates per bucket are quantize {TRT-LLM, CuTe DSL} x GEMM {CUTLASS, CuTe DSL}; on SM107 +# every TP2 projection at 1..16 rows runs fastest on this pair, cold and L2-hot (graph us, cold: 4.7..10.3 against +# 5.0..11.5 with the TRT-LLM quantize and 9.7..58 with CUTLASS), and all four candidates give the same bits (quantize +# bytes and scales, GEMM outputs). A constant is rank-agreed and start-independent by construction. +# False returns these projections to the AutoTuner. +MXFP8_DECODE_QUANT_TACTIC = 0 # MXFP8QuantizeRunner.CUTE_DSL +MXFP8_DECODE_GEMM_TACTIC = 0 # FlashInferMXFP8GemmRunner.CUTE_DSL +_ALL_ROWS = frozenset(range(1, TINY_MAX_TOKENS + 1)) +_GEMV_PROJECTIONS = ( + QKV_SPARSE, + QKV, + O_PROJ, + SHARED_GATE_UP, + SHARED_DOWN, + DENSE_GATE_UP, + DENSE_DOWN, +) +# B2: the same GEMVs at 9..16 rows (MMA width 16, the kernel's own wider configuration; contiguous k-tiles above 8 +# rows): standalone graph us at 16 rows 2.3..8.9 against 4.8..10.6 for the MX-TAC stock pair (b4p1 791712). Batch 8: +# at 17..32 rows (MMA width 32; rule 1b, a different K order than the generic GEMM). GEMM-AR standalone 794929 (iso +# latency vs the FlashInfer fallback on the same MXFP8 bytes), us per call at 24 / 32 rows: QKV-sparse -2.70 / -2.75, +# QKV -2.46 / -2.77, shared gate_up -3.12 / -3.12, shared down -0.94 / -1.01; dense gate_up +1.18 / +0.91, dense down +# +0.58 / +0.13 (the GEMV loses) and o_proj +0.11 / -0.18 (flat) keep 16 rows. Each projection is raised only as far +# as the kernel serves it (served_max_rows): the bf16 GEMV's in-kernel quantize declines dense gate_up / down past 8. +# MX_GEMV_B4 False: 9..32 rows decline to MX-TAC; B8_MX_GEMV_32 False: every GEMV to 16 rows (17..32-row MXFP8 +# inputs on FlashInfer's GEMM, exact). +MX_GEMV_B4 = True +B8_MX_GEMV_32 = True +GEMV_MAX_TOKENS = ( + (DECODE_MAX_TOKENS if B8_MX_GEMV_32 else B4_MAX_TOKENS) if MX_GEMV_B4 else TINY_MAX_TOKENS +) +_GEMV_B4_ROWS_ONLY = frozenset({DENSE_GATE_UP, DENSE_DOWN, O_PROJ}) + + +def _gemv_max(nk) -> int: + return min(GEMV_MAX_TOKENS, B4_MAX_TOKENS) if nk in _GEMV_B4_ROWS_ONLY else GEMV_MAX_TOKENS + + +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + { + nk: frozenset(range(1, served_max_rows(nk, False, TINY_MAX_TOKENS + 1, _gemv_max(nk)) + 1)) + for nk in _GEMV_PROJECTIONS + } + if MX_GEMV_TP2 + else {} +) +# W-G1 (agg wide decode, rule 1b): the MXFP8 GEMV also at 33..48 rows (MMA width 48 / 64) for the projections it wins +# on, after the decode-graph MXFP8 quantize the generic GEMM path runs anyway (the same activation bytes): standalone +# graph us per call at 48 rows, GEMV against the MX-TAC FlashInfer CuTe DSL GEMM on the same MXFP8 input, QKV-sparse +# 5.87 vs 10.82, QKV 5.53 vs 10.76, shared gate_up 4.52 vs 10.51, shared down 2.84 vs 4.74. o_proj and the dense MLP +# keep their rows (flat or losing at 24..48). A different K order than the GEMM (rule 1b). False: 33..48 rows back on +# the generic GEMM, exact. +WIDE_MX_GEMV_48 = True +WIDE_GEMV_MAX_TOKENS = 48 # capacity: fixed: the MXFP8 GEMV's widest build +_GEMV_W48_PROJECTIONS = (QKV_SPARSE, QKV, SHARED_GATE_UP, SHARED_DOWN) +_GEMV_W48_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + { + nk: frozenset( + range( + GEMV_MAX_TOKENS + 1, + served_max_rows(nk, False, GEMV_MAX_TOKENS + 1, WIDE_GEMV_MAX_TOKENS) + 1, + ) + ) + for nk in _GEMV_W48_PROJECTIONS + } + if MX_GEMV_TP2 and WIDE_MX_GEMV_48 and GEMV_MAX_TOKENS == DECODE_MAX_TOKENS + else {} +) +for _nk, _rows_w48 in _GEMV_W48_ROWS.items(): + MX_GEMV_ROWS[_nk] = MX_GEMV_ROWS.get(_nk, frozenset()) | _rows_w48 +MX_GEMV_BF16_MAX_TOKENS: Dict[Tuple[int, int], int] = ( + {nk: served_max_rows(nk, True, TINY_MAX_TOKENS + 1, _gemv_max(nk)) for nk in _GEMV_PROJECTIONS} + if MX_GEMV_TP2 + else {} +) +# W-G1: a bf16 activation (no MXFP8 hand-off) on rows the MXFP8 GEMV serves but the bf16 GEMV's in-kernel quantize does +# not (QKV / shared gate_up at 17..32 without the QKV hand-off, all four at 33..48) is quantized by the decode-graph +# MXFP8 quantize (the generic GEMM path's own) and runs the GEMV instead of that GEMM. +MX_GEMV_QUANT_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + { + nk: frozenset(r for r in MX_GEMV_ROWS.get(nk, ()) if r > MX_GEMV_BF16_MAX_TOKENS.get(nk, 0)) + for nk in _GEMV_W48_PROJECTIONS + } + if MX_GEMV_TP2 and WIDE_MX_GEMV_48 + else {} +) + +Key = Hashable + + +@dataclass(frozen=True) +class RowRange: + """The shapes one fast path takes, and the catalog entry whose cells certify each of them. + + ``keys`` are the covered shapes: row counts for most paths, (query rows per request, requests) for the attention + kernels, a variant name for thresholds. ``certifies`` maps a cell's spec of ``entry`` to the keys that cell + drives (empty for a cell of another variant). + """ + + path: str + keys: FrozenSet[Key] + entry: OpWrapper + certifies: Callable[[Mapping[str, Any]], Iterable[Key]] + + def covers(self, key: Key) -> bool: + return key in self.keys + + def certified(self) -> FrozenSet[Key]: + return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + + +def _rows(lo: int, hi: int) -> FrozenSet[int]: + return frozenset(range(lo, hi + 1)) + + +def _spec_rows(field: str, **match: Any) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + """A cell's ``spec[field]`` when the cell's other fields match ``match``.""" + + def certifies(spec: Mapping[str, Any]) -> Iterable[Key]: + if any(spec.get(name) != value for name, value in match.items()): + return () + return (spec[field],) + + return certifies + + +def _const(*keys: Key) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + return lambda spec: keys + + +def _sparse_decode_producer(spec: Mapping[str, Any]) -> Iterable[Key]: + if not spec["identity"]: + return () + rows = spec["decode_query_len"] * len(spec["seq_lens"]) + mq = rows >= SPARSE_DECODE_MQ_MIN_TOKENS and SPARSE_DECODE_KV_HEADS == 1 + return ("multi_query" if mq else "per_row",) + + +def index_score_schedule(requests: int) -> Tuple[int, bool]: + """The index scorer's (CTAs per request, early issue) at this many requests (a covered request count).""" + return INDEX_SCORE_SPLIT_K[requests], requests <= INDEX_SCORE_EARLY_MAX_REQUESTS + + +def _index_score_requests(spec: Mapping[str, Any]) -> Iterable[Key]: + """A cell certifies its request count when it drives the TP2 heads with that count's schedule.""" + requests = len(spec["seq_lens"]) + if spec.get("heads") != 2 or requests not in INDEX_SCORE_SPLIT_K: + return () + if (spec["split_k"], spec["early"]) != index_score_schedule(requests): + return () + return (requests,) + + +def _long_select_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The key widths a cell's rows reach: every width up to its blocks' when rows of mixed lengths are drawn.""" + blocks = spec["num_blocks"] + widest = 2 if blocks <= 2048 else 4 if blocks <= 4096 else 9 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_KEY_WIDTHS if w <= widest) + return (widest,) + + +def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The 384-thread select's key widths a cell's rows reach (every width up to its blocks' for mixed rows).""" + blocks = spec["num_blocks"] + widest = 4 if blocks <= 4 * 384 else 11 if blocks <= 11 * 384 else 24 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_DSL_KEY_WIDTHS if w <= widest) + return (widest,) + + +EMBED_FIRSTNORM = RowRange( + "embed_firstnorm", + _rows(1, EMBED_FIRSTNORM_MAX_TOKENS), + m3_embedding_firstnorm, + _spec_rows("rows"), +) +# One request decoding one token on an MXFP8 layer input: QKV GEMV + attention + o_proj in one launch. +# The identity-selection Triton producer + merge (every request within 16 blocks), certified at the TP2 heads on SM107. +SPARSE_DECODE_PRODUCER = RowRange( + _sparse_decode_kernel.PRODUCER_PATH, + _sparse_decode_kernel.PRODUCER_KEYS, + m3_sparse_decode, + _sparse_decode_producer, +) +# The long-context sparse decode at the TP2 heads (32 query / 2 KV heads): attention, split-K merge and the MXFP8 +# o_proj hand-off in one launch, for every decode step with a selected table. "mxfp8" is a step of at most +# MergeMxfp8.MAX_ROWS rows (the hand-off is written and published to m3_prequant.SLOT), "plain" a larger one. +# frozenset() routes every step back to the Triton producer + merge. +SPARSE_DECODE_LC = RowRange( + _sparse_decode_kernel.LC_PATH, + _sparse_decode_kernel.LC_KEYS, + m3_sparse_decode_lc, + lambda spec: ("mxfp8" if spec["mxfp8"] else "plain",), +) +# N1 at TP2: SPARSE_DECODE_LC records each step's selected block ids per sparse layer and, before its grid wait (while +# the select ranks), pulls into L2 the K/V pages + scales of the blocks the SAME layer selected one step back (65-68% of +# a layer's selected blocks recur one step later: MSA probe 778226). Same layer, not the next one as at TP4: the MoE +# between two sparse layers streams ~250 MB through the 126 MB L2, so an earlier issue would be evicted. Hints only; +# up to N1_MAX_ROWS query rows. False:. +SPARSE_DECODE_LC_PREFETCH_SELF_KV = True +N1_MAX_ROWS = 128 # capacity: fixed: the N1 prefetch's row bound +# Every token count of a sparse layer on an NVFP4 KV cache at the TP2 heads (prefill and decode): the producer takes the +# geometry, not a row range. frozenset() routes every sparse layer back to the native producer. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp2"}), m3_producer, _const("nvfp4_tp2")) +INDEX_SCORE = RowRange( + "m3_index_score", + _rows(1, INDEX_SCORE_MAX_REQUESTS), + m3_index_score, + _index_score_requests, +) +LONG_SELECT = RowRange( + "long_select", + frozenset(LONG_SELECT_KEY_WIDTHS) if LONG_SELECT_MAX_BLOCKS else frozenset(), + m3_select_long, + _long_select_widths, +) +LONG_SELECT_DSL = RowRange( + "long_select_dsl", + frozenset(LONG_SELECT_DSL_KEY_WIDTHS) + if LONG_SELECT_MAX_BLOCKS and SELECT_LONG_DSL + else frozenset(), + m3_select_long_dsl, + _long_select_dsl_widths, +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, PLAN_MC_AR_MAX_TOKENS) if MC_AR_TP2 else frozenset(), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=0), +) +MC_FINALIZE = RowRange( + "mc_finalize_allreduce_norm", + _rows(MC_FINALIZE_MIN_TOKENS, PLAN_MC_AR_MAX_TOKENS if B4_MC_FINALIZE else TINY_MAX_TOKENS), + m3_mc_finalize_allreduce_norm, + _spec_rows("m"), +) +ATTN_PUSH = RowRange( + "attn_push", + _rows(1, ATTN_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +DENSE_PUSH = RowRange( + "dense_push", + _rows(1, DENSE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +# The router of an unfused MoE of 9..48 rows (above the MoE engine's 8), a generation step's or a context step's: the +# TF32 tensor-core router (minimax_m3_router_tc, 16-CTA clusters, one reduction order; the arithmetic the built-in gate +# takes on SM107 at 9..16 generation rows), with the routed experts' NVFP4 activation from the same launch, as the agg +# canon routes these rows (2.9-3.5 us/call against ~9.8 us for the cuBLAS TF32 split-K router; closer to the fp32 +# product: max |dlogit| 1.7e-6 vs 6.7e-5). Devices that cannot co-schedule 16-CTA clusters keep the fp32 router. +ROUTER_TC_GEN_MIN_ROWS = 9 # capacity: fixed: past the MoE engine's 8 rows +# Up to 48 rows, as the agg canon runs it (v1.8.3, the router GEMV keys gone): the MoE of a 9..48-row step routes on +# it, with the routed experts' NVFP4 activation from the same kernel (the decode layer's MC finalize, the generic +# layer). GB300 tp4, behind a PDL producer, with the NVFP4 activation: 3.20 / 3.82 / 7.36 us at 8 / 12 / 48 tokens. +ROUTER_TC_GEN_MAX_ROWS = 48 # capacity: fixed: the router kernel's widest build +ROUTER_TC_GEN = RowRange( + "m3_router_tc", + _rows(ROUTER_TC_GEN_MIN_ROWS, ROUTER_TC_GEN_MAX_ROWS), + m3_router_tc, + _spec_rows("m"), # every cell checks the logits, with or without the NVFP4 activation +) +MOE_PUSH = RowRange( + "moe_push", + _rows(1, MOE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MOE_ENGINE = RowRange( + "m3_moe_engine", _rows(1, MOE_ENGINE_MAX_TOKENS), m3_moe_engine, _spec_rows("tokens") +) +# The GEMM paths are keyed (projection (N, K), rows): each projection runs its own split-K / ring / MMA +# configuration, so a row count is certified per projection. +MX_GEMV = RowRange( + "m3_mx_gemv", + frozenset((nk, m) for nk, rows in MX_GEMV_ROWS.items() for m in rows), + m3_mx_gemv, + lambda spec: ((spec["nk"], spec["m"]),), +) +MX_GEMV_BF16 = RowRange( + "m3_mx_gemv_bf16", + frozenset((nk, m) for nk, hi in MX_GEMV_BF16_MAX_TOKENS.items() for m in _rows(1, hi)), + m3_mx_gemv_bf16, + lambda spec: ((spec["nk"], spec["m"]),), +) +# (draft length, requests) of a greedy verify step. +SHARD_ACCEPT = RowRange( + "m3_shard_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SHARD_ACCEPT_MAX_REQUESTS + 1)), + m3_shard_accept, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) + +# Rows of a fused-draft-loop all-reduce boundary (drafter hidden 6144) on the multicast all-reduce. +DRAFT_MC_AR = RowRange( + "m3_draft_mc_allreduce_norm", + _rows(1, DRAFT_MC_AR_MAX_ROWS), + m3_mc_allreduce_norm, + lambda spec: (spec["m"],) if spec.get("drafter") else (), +) + +# (draft length, requests) of a sampled verify step. +SAMPLED_ACCEPT = RowRange( + "m3_sampled_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SAMPLED_VERIFY_MAX_REQUESTS + 1)), + m3_sampled_accept, + lambda spec: ((SHARD_ACCEPT_DRAFT_LEN, spec["nreq"]),), +) + + +# (query rows per request, requests) of a decode step whose dense-layer attention runs on PrimTS split-KV decode. +FULLCTX_DECODE = RowRange( + "m3_fullctx_decode", + frozenset(FULLCTX_DECODE_RECIPES) if FULLCTX_DECODE_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + +# (query rows per request, requests) of a generation-only drafter pass that runs on the full-context decode. +DRAFTER_FULLCTX = RowRange( + "m3_drafter_fullctx", + frozenset(FULLCTX_DECODE_RECIPES) if DRAFTER_FULLCTX_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + + +def _verify_rows(max_rows: int) -> FrozenSet[int]: + """The verify rows of 1..8 requests at the shard-accept draft length, up to ``max_rows``.""" + return frozenset( + n * (SHARD_ACCEPT_DRAFT_LEN + 1) + for n in range(1, 9) + if n * (SHARD_ACCEPT_DRAFT_LEN + 1) <= max_rows + ) + + +# Requests of a fused-verify step whose first drafter inputs are the fused ones; certified by the chunk-norm cells +# at those requests' verify rows (the norm is the only kernel of its own; the gather rows and the fc are in-tree). +FUSED_FIRST_INPUTS = RowRange( + "m3_fused_first_inputs", + _rows(1, FUSED_FIRST_INPUTS_MAX_REQUESTS), + m3_chunk_norm, + lambda spec: ( + (spec["m"] // (SHARD_ACCEPT_DRAFT_LEN + 1),) + if spec["m"] % (SHARD_ACCEPT_DRAFT_LEN + 1) == 0 + else () + ), +) +# Verify rows of a fused-verify step whose Eagle3 captures apply the fc_norm. +CAPTURE_NORM = RowRange( + "m3_capture_norm", _verify_rows(CAPTURE_NORM_MAX_ROWS), m3_capture_norm, _spec_rows("m") +) +# Verify rows of a first draft pass whose fc_norm runs as the chunk norm. +CHUNK_NORM = RowRange( + "m3_chunk_norm", _verify_rows(CHUNK_NORM_MAX_ROWS), m3_chunk_norm, _spec_rows("m") +) + +# ((N, K), rows) of a drafter BF16 linear on its low-M kernel. +DRAFT_LOWM = RowRange( + "m3_draft_lowm", + frozenset((nk, m) for nk, rows in DRAFT_LOWM_ROWS.items() for m in rows), + m3_draft_lowm, + lambda spec: ((spec["nk"], spec["m"]),), +) +# Verify rows of a step whose target logits come from the die-placed head (in-tree GEMM halves per die). +UGPU_HEAD = RowRange("m3_ugpu_head", frozenset(UGPU_HEAD_ROWS), m3_ugpu_head, _spec_rows("m")) + +# (draft length, requests) of a fused draft loop, and the rows of its passes' draft input: the first pass's +# requests x (draft length + 1), one per request on the later passes. +DRAFT_STEP = RowRange( + "m3_draft_step", + frozenset((DRAFT_STEP_DRAFT_LEN, n) for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1)), + m3_draft_step, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) +DRAFT_INPUT = RowRange( + "m3_draft_input", + frozenset( + rows + for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1) + for rows in (n, n * (DRAFT_STEP_DRAFT_LEN + 1)) + ), + m3_draft_input, + _spec_rows("m"), +) + +FAST_PATHS: Tuple[RowRange, ...] = ( + EMBED_FIRSTNORM, + FULLCTX_DECODE, + DRAFTER_FULLCTX, + SPARSE_DECODE_PRODUCER, + SPARSE_DECODE_LC, + PRODUCER, + INDEX_SCORE, + LONG_SELECT, + LONG_SELECT_DSL, + MC_ALLREDUCE, + MC_FINALIZE, + ATTN_PUSH, + DENSE_PUSH, + MOE_PUSH, + MOE_ENGINE, + ROUTER_TC_GEN, + MX_GEMV, + MX_GEMV_BF16, + SHARD_ACCEPT, + SAMPLED_ACCEPT, + DRAFT_MC_AR, + UGPU_HEAD, + DRAFT_STEP, + DRAFT_INPUT, + DRAFT_LOWM, + FUSED_FIRST_INPUTS, + CAPTURE_NORM, + CHUNK_NORM, +) + +# -------------------------------------------------------------------------------------------------------------------- +# Selection: pure functions of the step shape +# -------------------------------------------------------------------------------------------------------------------- + + +def plan_for(num_tokens: int) -> DecodePlan: + for plan in DECODE_PLANS: + if num_tokens <= plan.max_tokens: + return plan + raise AssertionError(f"no decode plan covers {num_tokens} tokens") + + +def is_decode_step(num_tokens: int) -> bool: + """Whether a pure generation step of ``num_tokens`` rows goes to the decode target.""" + return num_tokens <= DECODE_MAX_TOKENS + + +def moe_path(num_tokens: int) -> Optional[str]: + """The MoE kernel of a decode step: the engine, the unfinalized MoE + finalize all-reduce, or None (generic).""" + if MOE_ENGINE.covers(num_tokens): + return MOE_ENGINE.path + if MC_FINALIZE.covers(num_tokens): + return MC_FINALIZE.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# The dense layers' swigluoai activation as this target's C++ op (ops/gemm/m3_dense_swiglu, +# minimaxM3Sm107Tp2/swigluOai.cu, LongAttn c267035d2f: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, +# beta 1), PDL-launched with an early release so the down GEMV stages its weights during it; the down GEMV reads the +# activation only after its own grid wait. Indexer kdev17 783285 (TP2 chain): -0.69 us per dense layer vs the Triton +# kernel; LongAttn gate a 783511 (oracle 23/23). False runs the Triton kernel. +DENSE_SWIGLU_CPP_PATH = "m3_dense_swiglu" +# FlashInfer's CuTe DSL MXFP8 GEMM on the producer's MXFP8 activation (rows the GEMV does not serve). +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +# The MXFP8 quantize (tuned backend) and FlashInfer's CuTe DSL GEMM: every row the decode kernels do not serve. +GENERIC_GEMM_PATH = "mxfp8_gemm" + + +def linear_path(n: int, k: int, num_tokens: int, prequant: bool, low_m_gemm: bool) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. ``low_m_gemm`` is accepted for the call sites' + symmetry: no low-M GEMM configuration is certified at TP2, so it selects nothing. + """ + if prequant: + if num_tokens in MX_GEMV_ROWS.get((n, k), ()): + return MX_GEMV_PATH + return PREQUANT_GEMM_PATH + if num_tokens <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0): + return MX_GEMV_BF16_PATH + if num_tokens in MX_GEMV_QUANT_ROWS.get((n, k), ()): + # W-G1: the decode-graph MXFP8 quantize, then the MXFP8 GEMV. + return MX_GEMV_PATH + return GENERIC_GEMM_PATH + + +# -------------------------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A covered decode shape could not take its fast path under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() + + +def set_require(require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def missed(path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if _require[0]: + raise FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " + f"{reason}" + ) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def taken(path: str, rows: Optional[int], what: str = "") -> None: + """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line + every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" + at = "" if rows is None else f" at {rows} rows" + detail = f" ({what})" if what else "" + _log_once( + logger.info, + path, + f"taken: {rows}" + (f": {what}" if what else ""), + f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", + ) + + +def long_select_covers(num_blocks: int) -> bool: + """Whether ``LONG_SELECT`` takes rows of ``num_blocks`` KV blocks (the engine's per-row block width).""" + return LONG_SELECT_MIN_BLOCKS <= num_blocks <= LONG_SELECT_MAX_BLOCKS + + +def long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: + """The (PDL launch, early trigger, trigger phase) keys of the 384-thread select a decode step can launch, compiled + before warmup; none when ``LONG_SELECT_DSL`` covers nothing.""" + if not LONG_SELECT_DSL.keys: + return () + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + pdl = get_env_enable_pdl() + keys = [(pdl, False, SELECT_LONG_TRIGGER_PHASE)] + if SELECT_EARLY_MAX_ROWS: + keys.append((pdl, True, 0)) + return tuple(keys) + + +def long_select_fallback(num_blocks: int, rows: int) -> None: + """The in-tree selector is about to run instead of the long-row select: a miss when ``LONG_SELECT`` covers + ``num_blocks`` (the kernel declined operands its cells certify; raises under require), uncovered otherwise (rows of + at most 128 blocks, or wider than ``LONG_SELECT_MAX_BLOCKS``: the in-tree selector serves them in every mode).""" + if long_select_covers(num_blocks): + missed( + LONG_SELECT.path, + rows, + f"the long-row select declined {num_blocks}-block scores it covers", + ) + else: + span = f"{LONG_SELECT_MIN_BLOCKS}..{LONG_SELECT_MAX_BLOCKS}" + uncovered( + LONG_SELECT.path, f"{num_blocks} KV blocks, outside {span}: the built-in selector" + ) + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +# Kernel modules report from inside their calls through kernels/m3_report: the same policy. +m3_report.install( + on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed +) + + +# -------------------------------------------------------------------------------------------------------------------- +# Target constants +# -------------------------------------------------------------------------------------------------------------------- + +# The generic boundary takes the low-latency multicast all-reduce + residual + RMSNorm up to this many tokens; with +# the shared + routed add folded in, up to its own maximum: the decode boundaries' certified rows. +MC_AR_MAX_TOKENS = GENERIC_MC_AR_MAX_TOKENS +MC_AR_FOLDED_MAX_TOKENS = GENERIC_MC_AR_FOLDED_MAX_TOKENS +# The shared expert folds into the routed engine from one token (dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# SMs left to the concurrent kernels while a persistent routed engine runs. +MOE_ENGINE_RESERVED_SMS = 24 # capacity: fixed: an SM split, not a size +# Id of the target's final hidden state in Eagle3's ``layers_to_capture``. +FINAL_HIDDEN_LAYER_ID = -1 # capacity: fixed: an Eagle3 layer id, not a size +# The trtllm-gen MoE runner's routing kind for MiniMax-M3 (sigmoid + correction bias, renormalized, scaled) and its +# gated activation (SwiGLU with the per-expert alpha / beta / clamp operands). +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) +INDEX_SCORE_PATH = INDEX_SCORE.path + +# -------------------------------------------------------------------------------------------------------------------- +# Streaming: also emit a response on decoding iterations 2..FIRST_CONTENT_EMIT_ITERS (RuntimeHooks), so a request's +# first content leaves one step after its first token. 1: the stock schedule. +FIRST_CONTENT_EMIT_ITERS = 2 + +# The runtime-value dump (TRTLLM_MODELING_V2_VALUE_DUMP=1): the metadata of every fast-path and entry decision a step +# takes, recorded on eager steps and captures (replay runs no Python), printed once at exit. +# -------------------------------------------------------------------------------------------------------------------- + +_VALUE_DUMP = os.environ.get("TRTLLM_MODELING_V2_VALUE_DUMP") == "1" +_VALUES: set = set() + + +def _record(*meta: Hashable) -> None: + if _VALUE_DUMP: + _VALUES.add(meta) + + +def _print_values() -> None: + if _VALUES: + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "0")) + for meta in sorted(_VALUES, key=repr): + print(f"[modeling_v2 value dump] rank {rank}: {meta}", flush=True) + + +if _VALUE_DUMP: + atexit.register(_print_values) + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for (it is 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'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs +# -------------------------------------------------------------------------------------------------------------------- + + +def _mxfp8_linear( + core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + _record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +_PINNED_RUNNERS: dict = {} + + +def _pinned_runner(cls, dtype): + runner = _PINNED_RUNNERS.get((cls, dtype)) + if runner is None: + runner = _PINNED_RUNNERS[(cls, dtype)] = cls(dtype) + return runner + + +def _pinned_mxfp8_gemm(act, act_scale, weight, weight_scale, out_dtype): + """The decode-graph MXFP8 GEMM at the target's fixed tactic (``MXFP8_DECODE_GEMM_TACTIC``), no AutoTuner: the + same choice on every rank and every start.""" + _record("pinned_mxfp8_gemm", tuple(weight.shape), act.shape[0], MXFP8_DECODE_GEMM_TACTIC) + runner = _pinned_runner(FlashInferMXFP8GemmRunner, out_dtype) + return runner.forward([act, act_scale, weight, weight_scale], tactic=MXFP8_DECODE_GEMM_TACTIC) + + +def _pinned_mxfp8_quantize(x): + """The decode-graph MXFP8 activation quantize at the target's fixed tactic (``MXFP8_DECODE_QUANT_TACTIC``).""" + _record("pinned_mxfp8_quantize", x.shape[0], MXFP8_DECODE_QUANT_TACTIC) + return _pinned_runner(MXFP8QuantizeRunner, x.dtype).forward( + [x], tactic=MXFP8_DECODE_QUANT_TACTIC + ) + + +_MX_TAC_ENGAGED: set = set() + + +def _note_mx_tac(kind: str, n: int, k: int, rows: int) -> None: + """MX-TAC's engagement line, once per (activation kind, projection, rows).""" + key = (kind, n, k, rows) + if key not in _MX_TAC_ENGAGED: + _MX_TAC_ENGAGED.add(key) + logger.info( + f"[M3 sm107 tp2] MX-TAC engaged: pinned MXFP8 decision (quantize tactic " + f"{MXFP8_DECODE_QUANT_TACTIC}, GEMM tactic {MXFP8_DECODE_GEMM_TACTIC}) " + f"for the [{n}, {k}] projection at {rows} rows, {kind} activation" + ) + + +def _gemv_variants(n: int, k: int, push: Optional[str], world: int): + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch routes to an [N, K] + projection's GEMV: an MXFP8 activation on the rows ``MX_GEMV_ROWS`` gives it, a bf16 one up to + ``MX_GEMV_BF16_MAX_TOKENS``; each pushing when the projection's boundary push covers the rows. ``push``: "attn" + (the o_proj into the attention boundary), "dense" (the dense MLP's down projection into the post-MLP boundary) + or "moe" (the shared expert's down projection into the MoE boundary).""" + push_rows = {"attn": ATTN_PUSH.keys, "dense": DENSE_PUSH.keys, "moe": MOE_PUSH.keys}.get( + push, frozenset() + ) + variants = [] + max_rows = max([DECODE_MAX_TOKENS, *(max(r) for r in MX_GEMV_ROWS.values() if r)]) + for rows in range(1, max_rows + 1): + for x_bf16, served in ( + (False, rows in MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + bucket = next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + variants.append((rows, x_bf16, (world, bucket, AR_PUSH_MODE))) + return variants + + +def _uncompiled_note() -> str: + key = mx_gemv.last_uncompiled + return "" if key is None else f" (its kernel {key} was not compiled before warmup)" + + +def _linear( + core, x: torch.Tensor, wk: str, sk: str, low_m_gemm: bool, push: Optional[tuple] = None +) -> Tuple[torch.Tensor, bool]: + """One MXFP8 projection ``x @ w[wk]^T``; returns (output, whether its GEMV pushed into ``push``). + + Outside decode CUDA graphs: the engine-tuned MXFP8 quantize + GEMM. Inside them the projection's (N, K) and the + step's rows pick one path (``linear_path``): + + 1. an activation its producer already quantized to MXFP8 (``m3_prequant.SLOT``, taken here): the tcgen05 decode + GEMV on the rows it serves, else the decode-graph MXFP8 GEMM on that activation; + 2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; + 3. W-G1: a bf16 activation on the rows the MXFP8 GEMV serves but the bf16 GEMV does not + (``MX_GEMV_QUANT_ROWS``): the decode-graph MXFP8 quantize, then the MXFP8 GEMV; + 4. otherwise the decode-graph MXFP8 quantize + GEMM. + + MX-TAC: on a decode graph whose MXFP8 projections the engine tunes (``tune_decode_graph_backends``), the GEMM + (and, on a bf16 activation, the quantize) the GEMVs do not serve take one fixed decision + (``MXFP8_DECODE_QUANT_TACTIC`` / ``MXFP8_DECODE_GEMM_TACTIC``: FlashInfer's CuTe DSL quantize and CuTe DSL GEMM) + instead of the AutoTuner's per-start choice (whose L2-hot profiling left c2 bimodal at +-6%). On SM107 every TP2 + projection at 1..16 rows runs fastest on this pair, cold and L2-hot (graph us, cold: 4.7..10.3 against 5.0..11.5 + with the TRT-LLM quantize and 9.7..58 with CUTLASS), and all four candidates give the same bits (quantize bytes + and scales, GEMM outputs). A constant is rank-agreed and start-independent by construction. + + A kernel declining rows its path covers is a miss; the chain then continues down the list. ``push`` (an + ``m3_mx_gemv`` ``ar`` tuple) is the boundary all-reduce exchange the GEMV writes its rows into. ``low_m_gemm`` + marks the projections the fused low-M MXFP8 GEMM may take (``LOW_M_GEMM_MAX_TOKENS``: no TP2 configuration is + built, so ``linear_path`` never selects it). + """ + w, s = core.w[wk], core.w[sk] + x = x.contiguous() + mode = _mx_mode(core.mxfp8) + if not mode.decode_graph: + return _mxfp8_linear(core, x, w, s, mode), False + use_pdl = get_env_enable_pdl() + rows = x.shape[0] + n, k = w.shape + prequant = m3_prequant.SLOT.take(x) + if prequant is None and x.dtype != torch.bfloat16: + path = GENERIC_GEMM_PATH + else: + path = linear_path(n, k, rows, prequant is not None, low_m_gemm) + _record("linear", (n, k), rows, path, push is not None) + if prequant is not None: + act, act_scale = prequant + if path == MX_GEMV_PATH: + output = m3_mx_gemv(act, act_scale, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] mxfp8 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if path == PREQUANT_GEMM_PATH: + uncovered(MX_GEMV_PATH, f"the [{n}, {k}] projection at rows the GEMV does not serve") + if mode.tuned_quantize_gemm: + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("mxfp8", n, k, rows) + return output, False + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_PATH: + # W-G1 (MX_GEMV_QUANT_ROWS): no MXFP8 hand-off, so the decode-graph quantize the generic GEMM path would run + # (the same bytes), then the GEMV in place of that GEMM. + act, act_scale = _pinned_mxfp8_quantize(x) + output = m3_mx_gemv(act, act_scale, w, s, True, None, use_pdl, None) + if output is not None: + taken(path, rows, f"[{n}, {k}] mxfp8 quantize + GEMV") + return output, False + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's quantized activation" + + _uncompiled_note(), + ) + if mode.tuned_quantize_gemm: + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("mxfp8", n, k, rows) + return output, False + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16(x, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] bf16 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, rows, f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note() + ) + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + if mode.tuned_quantize_gemm and x.dtype == torch.bfloat16: + act, act_scale = _pinned_mxfp8_quantize(x) + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("bf16", n, k, rows) + return output, False + return _mxfp8_linear(core, x, w, s, mode), False + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer boundaries: all-reduce + residual + RMSNorm +# -------------------------------------------------------------------------------------------------------------------- + + +def _can_mc(x, folded: bool = False) -> bool: + """Whether the generic boundary takes the multicast all-reduce + residual + RMSNorm: the rows + ``MC_AR_MAX_TOKENS`` / ``MC_AR_FOLDED_MAX_TOKENS`` (with the shared + routed add folded in) give it.""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (MC_AR_FOLDED_MAX_TOKENS if folded else MC_AR_MAX_TOKENS) + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + +def _mc_allreduce_norm(core, x, x2, residual, norm_w, *, quant=False, prepushed=0, capture=None): + """The multicast all-reduce + optional second addend + residual + RMSNorm (``m3_mc_allreduce_norm``), or None + without a workspace or for operands the kernel does not take (the same on every rank). With ``quant`` the normed + rows are also published, MXFP8-quantized, to ``m3_prequant.SLOT`` for the next MXFP8 projection; ``capture`` = + (out, fc_norm weight, eps) also writes the layer's Eagle3 capture norm of the updated residual.""" + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + if workspace is None: + return None + _record( + "m3_mc_allreduce_norm", x.shape[0], x2 is not None, quant, prepushed, capture is not None + ) + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm_w, + core.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + capture=capture, + ) + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm.""" + _record("m3_allreduce", x.shape[0]) + return m3_allreduce( + x, + core.allreduce, + AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ), + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The MoE: router, routed experts (trtllm-gen NVFP4), shared expert +# -------------------------------------------------------------------------------------------------------------------- + + +@dataclass +class _MoeScalars: + """The routed experts' kernel scalars of one layer, derived from the checkpoint's per-tensor scales once the + weights are real (``MiniMaxM3Tp2ModelingV2._derive_moe``).""" + + fc31_input_scale: torch.Tensor + fc2_input_scale: torch.Tensor + fc31_alpha: torch.Tensor + fc2_alpha: torch.Tensor + fc31_scale_c: torch.Tensor + act_alpha: torch.Tensor + act_beta: torch.Tensor + act_clamp: torch.Tensor + + +def _router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The fp32 router logits of MoE layer ``i`` when the MoE runs unfused outside ``ROUTER_TC_GEN``'s rows (the decode + steps' engines route in-kernel): ``F.linear(float(h), w)``, the built-in gate's generation-step arithmetic; the + experts then quantize their bf16 input themselves.""" + _record("m3_fp32_linear", h.shape[0] if h.dim() == 2 else 0) + return m3_fp32_linear(h, core.w[f"l{i}_router"]) + + +def _decode_router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The router logits of a pure generation step on the generic layer: ``ROUTER_TC_GEN``'s rows take the TF32 + tensor-core router when the device co-schedules 16-CTA clusters (decided at prebuild); every other step the fp32 + router (``_router_logits``).""" + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + ROUTER_TC_GEN.covers(rows) + and core.router_tc_cluster16 + and h.dtype == torch.bfloat16 + and h.stride(1) == 1 + and (h.stride(0) * 2) % 16 == 0 + ): + taken(ROUTER_TC_GEN.path, rows) + _record("m3_router_tc", rows) + return m3_router_tc(h, core.w[f"l{i}_router"], None, get_env_enable_pdl())[0] + return _router_logits(core, i, h) + + +def _router_tc_fp4(core, i: int, h: torch.Tensor) -> Optional[tuple]: + """(router logits, NVFP4 activation, its scales) of an unfused MoE on ``ROUTER_TC_GEN``'s rows, a generation or a + context step: the TF32 tensor-core router with the routed experts' NVFP4 activation quantization fused in (the agg + canon's router on these rows); None elsewhere (the caller routes on its fp32 router).""" + rows = h.shape[0] if h.dim() == 2 else 0 + if not ( + ROUTER_TC_GEN.covers(rows) + and core.router_tc_cluster16 + and h.dtype == torch.bfloat16 + and h.stride(1) == 1 + and (h.stride(0) * 2) % 16 == 0 + ): + return None + taken(ROUTER_TC_GEN.path, rows, "with the NVFP4 activation") + _record("m3_router_tc", rows, True) + return tuple( + m3_router_tc( + h, core.w[f"l{i}_router"], core.moe[i].fc31_input_scale.reshape(1), get_env_enable_pdl() + ) + ) + + +def _decode_routed_unfinalized(core, i: int, h: torch.Tensor): + """The routed experts of a pure generation step, unfinalized (``[gemm2 out, expert scales, expanded->permuted + index]``): routed on the fused tensor-core router where it covers the rows, else on the router logits with the + experts quantizing their bf16 input.""" + fused = _router_tc_fp4(core, i, h) + if fused is not None: + logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, logits, False) + return _routed_experts(core, i, h, None, _decode_router_logits(core, i, h), False) + + +def _triton_finalize(rows: int, routed, shared: Optional[torch.Tensor]) -> torch.Tensor: + """The finalized routed output (+ ``shared``) from the runner's unfinalized parts via kernels/m3_moe_finalize + (``MOE_TRITON_FINALIZE``). Raises if the parts are not the trtllm-gen layout (the runner did not finalize them: + no fallback).""" + out = None + if isinstance(routed, (list, tuple)) and len(routed) == 3: + out = m3_moe_finalize(routed[0], routed[1], routed[2], shared) + if out is None: + raise RuntimeError( + f"the routed experts' unfinalized output at {rows} rows is not the trtllm-gen layout" + ) + taken( + "m3_moe_finalize", rows, "Triton finalize" + (" + shared add" if shared is not None else "") + ) + return out + + +def _triton_finalize_rows(core, h: torch.Tensor) -> bool: + """Whether a finalized MoE forward of ``h`` takes ``MOE_TRITON_FINALIZE``: at least its rows, in one runner call.""" + rows = h.shape[0] if h.dim() == 2 else 0 + return ( + MOE_TRITON_FINALIZE + and rows >= MOE_TRITON_FINALIZE_MIN_TOKENS + and rows <= core.moe_max_num_tokens + ) + + +def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): + """The trtllm-gen NVFP4 MoE runner over this rank's tensor-parallel slice of every expert, routing in-kernel. + + ``x_sf`` None means ``x`` is the bf16 activation, quantized here with the experts' global input scale. Returns + the finalized rows, or the unfinalized ``[gemm2 out, expert scales, expanded->permuted index]``. + """ + moe = core.moe[i] + if x_sf is None: + x, x_sf = fp4_quantize(x, moe.fc31_input_scale, 16, False, False) + rows = x.shape[0] + x_sf = x_sf.view(rows, -1).flatten() + _record("fp4_block_scale_moe_runner", rows, do_finalize) + outs = fp4_block_scale_moe_runner( + router_logits, + core.w[f"l{i}_router_bias"], + x, + x_sf.view(torch.float8_e4m3fn), + core.w[f"l{i}_e_w31"], + core.w[f"l{i}_e_s31"].view(torch.float8_e4m3fn), + None, + moe.act_alpha, + moe.act_beta, + moe.act_clamp, + core.w[f"l{i}_e_w2"], + core.w[f"l{i}_e_s2"].view(torch.float8_e4m3fn), + None, + moe.fc31_scale_c, + moe.fc31_alpha, + moe.fc2_alpha, + core.num_experts, + core.top_k, + None, + None, + core.moe_inter, + 0, + core.num_experts, + core.routed_scaling_factor, + ROUTING_MINIMAX2, + do_finalize, + act_type=ACT_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=core.max_num_tokens, + use_dp=False, + ) + return outs[0] if do_finalize else outs + + +def _split_chunk(rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +# -------------------------------------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------------------------------------- + + +def _nvfp4_layer(manager, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _producer_covered(core) -> bool: + """Whether ``PRODUCER`` covers this target's sparse layers: an NVFP4 KV cache at the TP2 heads (32 query, two KV + and two index heads). A covered layer that cannot take ``m3_producer`` is a miss.""" + return core.main_kv_is_nvfp4 and (core.heads, core.kv_heads, core.index_heads) == (32, 2, 2) + + +def _producer_layout(kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 pools are the layout ``m3_producer`` is built for: [pages, 2, 2, 128, 64] with + its scale pool [pages, 2, 2, 128, 8].""" + return ( + kv_data is not None + and kv_scale is not None + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 2, 128, 64) + and kv_data.stride()[2:] == (128 * 64, 64, 1) + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 2, 128, 8) + and kv_scale.stride()[2:] == (128 * 8, 8, 1) + and kv_scale.shape[0] == kv_data.shape[0] + ) + + +def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or data.shape[0] != scales.shape[0] + or tuple(data.shape[1:]) != (2, num_heads, 128, 64) + or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) + or data.dtype not in (torch.int8, torch.uint8) + or scales.dtype != torch.uint8 + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or cache.stride(1) < num_heads * cache.stride(2) + or cache.stride(0) < 2 * cache.stride(1) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + + +#: The step phases ("decode" / "prefill") whose producer engagement has been logged, once per process. +_PRODUCER_ENGAGED: set = set() + + +def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: + if covered: + missed(PRODUCER.path, rows, reason) + return None + + +def _sparse_qkv_producer( + core, i: int, packed, position_ids, kv_cache_manager, msa_idx_k_cache, msa_out_cache_loc +): + """The horizontal producer of a sparse layer: Q / index-Q (FP8) out, K / V / index-K written into the paged + caches (NVFP4 or FP8 main cache), for every step type (token-major). An NVFP4 cache at the TP2 heads takes + ``m3_producer`` (``PRODUCER``); a covered layer that cannot is a miss. Returns (q, index_q), or None for caches + neither kernel takes. + + ``kv_cache_manager``, ``msa_idx_k_cache`` (the per-layer index-K cache getter) and ``msa_out_cache_loc`` are the + calling target's ``step_args`` entries.""" + covered = _producer_covered(core) + rows = int(packed.shape[0]) + if packed.dtype != torch.bfloat16 or position_ids is None: + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + manager = kv_cache_manager + if manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = manager.get_buffers(i, kv_layout="HND") + index_k_cache = msa_idx_k_cache(i) + out_cache_loc = msa_out_cache_loc + layer_uses_nvfp4 = _nvfp4_layer(manager, i) + scale_buffers = ( + manager.get_block_scale_buffers(i, kv_layout="HND") if layer_uses_nvfp4 else None + ) + kv_quant_scale = core.inv_kv_scales[i] if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supported_main = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_fused_write_layout(buffers, scale_buffers, core.kv_heads) + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale.is_cuda + and kv_quant_scale.dtype == torch.float32 + and kv_quant_scale.numel() >= 3 + and kv_quant_scale.is_contiguous() + ) + else: + supported_main = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, core.kv_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= core.kv_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + cos_sin = core.rotary_cos_sin + supported_rope = ( + cos_sin.is_cuda + and cos_sin.dtype == torch.float32 + and cos_sin.is_contiguous() + and cos_sin.dim() == 3 + and tuple(cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main + or not supported_index + or not supported_rope + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < rows + ): + return _producer_miss( + covered and layer_uses_nvfp4, + rows, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the layout the " + "producer is built for", + ) + w = core.w + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + layer_uses_nvfp4 + and (core.heads, core.kv_heads, core.index_heads) == (32, 2, 2) + and _producer_layout(buffers, scale_buffers) + ): + _record("m3_producer", rows) + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + core.heads, + get_env_enable_pdl(), + PRODUCER_EARLY_TRIGGER, + PRODUCER_PREWAIT_STEP_INPUTS, + ) + phase = "decode" if rows <= DECODE_MAX_TOKENS else "prefill" + if phase not in _PRODUCER_ENGAGED: + _PRODUCER_ENGAGED.add(phase) + logger.info( + f"[M3 sm107 tp2] fast path engaged: m3_producer ({phase}, T={rows}, " + f"early_trigger={PRODUCER_EARLY_TRIGGER}, prewait_step_inputs={PRODUCER_PREWAIT_STEP_INPUTS})" + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool {tuple(scale_buffers.shape)} is not the " + "producer's layout", + ) + _record("m3_qkv_insert", rows, layer_uses_nvfp4) + q, index_q = m3_qkv_insert( + packed.contiguous(), + buffers.view(torch.uint8) if layer_uses_nvfp4 else buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.heads, + core.kv_heads, + core.index_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + ) + return q.flatten(1), index_q.flatten(1) + + +def _qk_norm_rope(core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool) -> torch.Tensor: + """The per-head Gemma RMSNorm of Q and K and the partial NeoX RoPE of a dense layer's packed [Q | K | V], in one + kernel: in place, or an FP8 copy (Q / K normed and roped, V cast) when ``out_fp8``.""" + qkv = qkv.contiguous() + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + args = ( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + 1.0, # factor: no YARN + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + _record("qk_norm_rope", qkv.shape[0], out_fp8) + if out_fp8: + return m3_qk_norm_rope_fp8(*args) + return fused_qk_norm_rope(*args) + + +def _dense_fp8_op_applies(qkv: torch.Tensor, position_ids) -> bool: + """Whether a dense layer's FP8 q/k/v come from this target's op (``minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8``, + else the BF16 fused_qk_norm_rope output cast to FP8).""" + return ( + DENSE_QKV_FP8_VIA_BF16 + and hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8") + and position_ids is not None + and qkv.dtype == torch.bfloat16 + ) + + +def _dense_qk_norm_rope_fp8( + core, i: int, qkv: torch.Tensor, position_ids, kv_dest: Optional[tuple] = None +) -> torch.Tensor: + """A dense layer's FP8 q/k/v: QK norm + RoPE with Q/K rounded to BF16 and then to E4M3 and V cast, the bits of the + BF16 fused_qk_norm_rope output cast to FP8 (K is stored in the FP8 cache), in one launch of this target's op + (``DENSE_QKV_FP8_VIA_BF16``); without the op, the BF16 kernel and the cast. + + With ``kv_dest`` = (the layer's HND pool view, the step's int32 cache slots, live tokens, a [tokens, q heads, 128] + FP8 buffer) the op also stores the live tokens' K and V rows into the pool (the bytes ``write_msa_phase_kv`` + stores, one launch instead of about ten) and every row's Q into the buffer; ``kv_dest`` is passed only where the + op is known to run (``_dense_fp8_op_applies``).""" + if not _dense_fp8_op_applies(qkv, position_ids): + return fp8_e4m3_saturate(_qk_norm_rope(core, i, qkv, position_ids, False)) + _record("m3_dense_qk_norm_rope_fp8", qkv.shape[0], kv_dest is not None) + args = ( + qkv.contiguous(), + core.heads, + core.kv_heads, + core.kv_heads, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, # NeoX RoPE + position_ids.reshape(-1).contiguous().to(torch.int32), + True, # Gemma norms + ) + if kv_dest is not None: + return torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv(*args, *kv_dest) + return torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + + +def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: + """The dense MLP's SwiGLU-OAI activation: this target's C++ op (``m3_dense_swiglu``, minimaxM3Sm107Tp2/ + swigluOai.cu: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, beta 1, PDL-launched with an early + release so the down GEMV stages its weights during it) at the rows its cells certify, in any graph or eager step; + otherwise the Triton kernel (``m3_silu_and_mul``). Indexer kdev17 (TP2 chain): -0.69 us per dense layer vs the + Triton kernel. Every decline is logged once.""" + path = DENSE_SWIGLU_CPP_PATH + rows = gu.shape[0] + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + params = (core.swiglu_alpha, core.swiglu_limit, core.swiglu_beta) + act = None + if rows > max_rows: + uncovered(path, f"{rows} rows, past its certified {max_rows}") + elif params != (DENSE_ALPHA, DENSE_LIMIT, DENSE_BETA): + uncovered(path, f"activation parameters {params}") + else: + act = m3_dense_swiglu( + gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + ) + if act is None: + uncovered(path, "the gate_up layout") + else: + _record("m3_dense_swiglu", rows) + taken(path, rows, "dense swigluoai") + if act is None: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return act + + +def _fp8_index_qk_norm_rope( + core, i: int, idx_qk: torch.Tensor, position_ids, msa_idx_k_cache, msa_out_cache_loc +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + _record("m3_fp8_indexer_qk_norm_rope", rows) + return m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer helpers both targets' bodies call: the Eagle3 capture hook and the norm after a layer +# -------------------------------------------------------------------------------------------------------------------- + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + +def _next_norm(core, i: int) -> torch.Tensor: + """The norm the boundary after layer ``i`` applies: the next layer's input norm, the final norm after the last.""" + return core.w[f"l{i + 1}_in_norm"] if i + 1 < core.num_layers else core.w["final_norm"] + + +# -------------------------------------------------------------------------------------------------------------------- +# Kernel prebuilds the core runs before the executor's warmup +# -------------------------------------------------------------------------------------------------------------------- + + +def _in_tree_scorer(): + """The in-tree CuTe DSL index scorer's runner (importing it registers ``cute_dsl_minimax_m3_index_decode_score``); + imported on first use, as the CuTe DSL ops are.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLMiniMaxM3IndexDecodeScoreRunner, + ) + + return CuteDSLMiniMaxM3IndexDecodeScoreRunner + + +def _dummy_scorer_operands(dtype, heads, qlen, dim, page, device): + """One request of two blocks: real-shaped operands that load a compiled scorer module.""" + rows, blocks = 1, 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros(rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view(rows, blocks) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + return idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + + +#: Mixed (context + generation) steps run their generation rows on the built-in sparse decode +#: (``MsaDecodeFmha._run_sparse``). Its Triton producer and merge specialize on the split-K factor (16 / 8 / 4 / 2 / 1 +#: by the row count on SM107), on the row count's divisibility and on the 16-byte alignment of the request rows (a +#: mixed step's generation requests follow its context requests, so its block-table and length rows start mid-buffer), +#: so ``_prebuild_builtin_sparse_decode`` launches them per decode row count up to this many rows, both alignments. +BUILTIN_SPARSE_DECODE_PREBUILD_MAX_ROWS = ( # capacity: fixed: the prebuilt built-in decode rows + 64 # capacity: fixed: the mixed steps' generation rows prebuilt +) + + +def _prebuild_builtin_sparse_decode(query_lens, num_heads, num_kv_heads, device) -> list: + """Compile (and load) the built-in sparse decode's Triton producer and merge for every generation row count a mixed + step can have, so no step compiles one after the server is ready. Each launch mirrors the operands + ``MsaDecodeFmha._run_sparse`` passes for an NVFP4 layer (FP8 q, the HND pool's K / V planes, the scale pool's + planes, the head-major top-k view, BF16 output) on dummy values, on the current stream. Returns the row counts.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode, + ) + + page, head_dim, topk = 128, 128, MSA_REQUIRED_TOPK + pages = topk + pool = torch.zeros( + pages, 2, num_kv_heads, page, head_dim // 2, dtype=torch.uint8, device=device + ) + scales = torch.zeros( + pages, 2, num_kv_heads, page, head_dim // 16, dtype=torch.uint8, device=device + ) + global_scales = torch.ones(2, 4, dtype=torch.float32, device=device) + done = [] + for qlen in sorted(set(query_lens)): + for rows in range(qlen, BUILTIN_SPARSE_DECODE_PREBUILD_MAX_ROWS + 1, qlen): + requests = rows // qlen + q = torch.zeros(rows, num_heads, head_dim, dtype=torch.bfloat16, device=device).to( + torch.float8_e4m3fn + ) + indexes = ( + torch.arange(topk, dtype=torch.int32, device=device) + .repeat(rows, num_kv_heads, 1) + .contiguous() + ) + # One leading context request: its rows are skipped, so the generation rows start mid-buffer. + block_rows = torch.arange(topk, dtype=torch.int32, device=device).repeat( + requests + 1, 1 + ) + lengths = torch.full((requests + 1,), topk * page, dtype=torch.int32, device=device) + output = torch.empty(rows, num_heads, head_dim, dtype=torch.bfloat16, device=device) + for first in (0, 1): + minimax_m3_sparse_attn_decode( + q, + pool[:, 0], + pool[:, 1], + indexes.permute(1, 0, 2), + block_rows[first : first + requests], + lengths[first : first + requests], + sm_scale=head_dim**-0.5, + output=output, + decode_query_len=qlen, + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=global_scales[0, :1], + v_global_scale=global_scales[1, :1], + ) + done.append(rows) + return done + + +# -------------------------------------------------------------------------------------------------------------------- +# The core +# -------------------------------------------------------------------------------------------------------------------- + +# Engine states (compiled kernel + workspace) per (device, token count[, shared fold]), shared by every layer: a step +# runs the layers one after another on one stream. +_ENGINE_STATES: Dict[tuple, object] = {} + + +class MiniMaxM3Tp2ModelingV2(DecoderModel): + """The MiniMax-M3 text model of this target: its weights, its runtime resources, and the per-step dispatch.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_runtime(model_config) + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + self.mapping = mapping + # Whether the 9..16-row generation router takes the TF32 tensor-core kernel (set by prebuild_kernels). + self.router_tc_cluster16 = False + self.vocab_size = cfg.vocab_size + d = _weights.MODEL_WEIGHTS.dims(self) + self.num_layers = d.num_layers + self.hidden = d.hidden + self.heads = d.heads + self.kv_heads = d.kv_heads + self.head_dim = d.head_dim + self.index_heads = d.index_heads + self.index_dim = d.index_dim + self.q_size = d.heads * d.head_dim + self.kv_size = d.kv_heads * d.head_dim + self.moe_inter = d.moe_inter + self.num_experts = d.num_experts + self.top_k = int(cfg.num_experts_per_tok) + self.routed_scaling_factor = float(getattr(cfg, "routed_scaling_factor", 1.0)) + self.rms_eps = float(cfg.rms_norm_eps) + self.swiglu_alpha = float(getattr(cfg, "swiglu_alpha", 1.702)) + self.swiglu_beta = 1.0 # the (up + 1) offset of SwiGLU-OAI + self.swiglu_limit = float(getattr(cfg, "swiglu_limit", 7.0)) + self.max_num_tokens = model_config.max_num_tokens + self.moe_max_num_tokens = model_config.moe_max_num_tokens + self.sparse_layers = frozenset(_sparse_layer_ids(cfg)) + self.moe_layers = frozenset(_moe_layer_ids(cfg)) + mode = model_config.quant_config.quant_mode + self.main_kv_is_fp8 = bool(mode is not None and mode.has_fp8_kv_cache()) + self.main_kv_is_nvfp4 = bool(mode is not None and mode.has_fp4_kv_cache()) + # The sparse layers' QKV and index projections: one five-way packed projection, or the QKV projection and a + # separate BF16 index projection (fuse_qkv_index_projection=False) run on two streams. + self.fused_projection = bool( + getattr(model_config.sparse_attention_config, "fuse_qkv_index_projection", False) + ) + + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # The engine reads these two off the module tree: the MXFP8 GEMM state it tunes, and the opt-in to the + # automatic FlashInfer MXFP8 path on decode graphs. + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + + # The built-in fused all-reduce + residual + RMSNorm of the layer boundaries the multicast kernel does not take. + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + + # Streams: the shared / routed expert overlap and, when the MoE chunks its input, the chunking overlap. + self.aux_stream_moe = torch.cuda.Stream() + chunked = self.moe_max_num_tokens < model_config.max_num_tokens * mapping.dp_size + self.aux_stream_moe_chunk = torch.cuda.Stream() if chunked else None + self.moe_events = {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers} + self.moe_chunk_events = { + i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers + } + # The unfused sparse layers' index branch overlaps the QKV branch on its own stream. + self.aux_stream_attn = None if self.fused_projection else torch.cuda.Stream() + self.attn_events = ( + {} + if self.fused_projection + else {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.sparse_layers} + ) + + # Partial NeoX RoPE (rotary_dim of head_dim): the cos/sin table every attention kernel reads. + rope = RopeParams.from_config(cfg) + self.rope_theta = rope.theta + self.rotary_dim = int(rope.dim) + self.rotary_cos_sin = rope.create_rope_const_params(interleave=False)[1].reshape( + rope.max_positions, 2, -1 + ) + + # One MSA attention backend per layer (dense and sparse; the step metadata is the backend's). + quant_config = model_config.get_quant_config() + self.attn: List[MiniMaxM3MsaSparseAttention] = [] + for i in range(self.num_layers): + backend = create_attention( + model_config.attn_backend, + i, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=None, + quant_config=quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + q_scaling=1.0, + sparse_params=model_config.sparse_attention_config.to_sparse_params( + pretrained_config=cfg, layer_idx=i + ), + ) + backend.update_quant_config(quant_config) + if not isinstance(backend, MiniMaxM3MsaSparseAttention): + raise ValueError( + f"layer {i}: expected the MiniMax-M3 MSA backend, got {type(backend).__name__}" + ) + if i not in self.sparse_layers: + # The KV split trtllm-gen dense decode runs at, per request count. + backend.dense_decode_max_seq_len_by_batch = DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + self.attn.append(backend) + + # Built once the weights are real (``post_load_weights``): the KV (Q, K, V) scales, the MoE scalars, the + # decode engines and workspaces, the two targets. + self.kv_scales: Dict[int, torch.Tensor] = {} + self.inv_kv_scales: Dict[int, torch.Tensor] = {} + self.moe: Dict[int, _MoeScalars] = {} + self.mx_alpha: Optional[torch.Tensor] = None + self.moe_engines: Dict[int, Dict[Tuple[int, bool, bool], Optional[object]]] = {} + self.embed_owners: Dict[tuple, DecodeEmbedding] = {} + # Each sparse layer's selected block ids per (KV head, query row, slot) at its last long-context decode, which + # that decode prefetches into L2 before its grid wait (``SPARSE_DECODE_LC_PREFETCH_SELF_KV``, up to + # ``N1_MAX_ROWS`` query rows). + self.n1_ids: Dict[int, torch.Tensor] = ( + { + i: torch.full( + (self.kv_heads, N1_MAX_ROWS, MSA_REQUIRED_TOPK), + -1, + dtype=torch.int32, + device=torch.device("cuda", torch.cuda.current_device()), + ) + for i in sorted(self.sparse_layers) + } + if SPARSE_DECODE_LC_PREFETCH_SELF_KV + else {} + ) + self._built_for: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def sm_scale(self, i: int) -> float: + return (self.head_dim**-0.5) / float(self.attn[i].q_scaling) + + # ---------------------------------------------------------------------------------------------- after the load + + def post_load_weights(self) -> None: + """Derive the kernels' state from the loaded weights, then build both targets and every decode kernel. + + Idempotent per set of weight storages: the loader's post-load walk reaches this module directly and through + the shell, and a second call on the same storages does nothing. + """ + identity = ( + self.w["embed"].data_ptr(), + self.w["l0_in_norm"].data_ptr(), + self.w["final_norm"].data_ptr(), + ) + if self._built_for == identity: + return + device = self.w["embed"].device + derived = nn.Module() + derived.register_buffer( + "mx_alpha", torch.ones(1, dtype=torch.float32, device=device), persistent=False + ) + for i in range(self.num_layers): + # No KV calibration in the checkpoint: unit (Q, K, V) scales and their reciprocals. + derived.register_buffer( + f"l{i}_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + derived.register_buffer( + f"l{i}_inv_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + for i in sorted(self.moe_layers): + for name, value in dataclasses.asdict(self._derive_moe(i)).items(): + derived.register_buffer(f"l{i}_{name}", value, persistent=False) + self.derived = derived + self.mx_alpha = derived.mx_alpha + self.kv_scales = {i: getattr(derived, f"l{i}_kv_scales") for i in range(self.num_layers)} + self.inv_kv_scales = { + i: getattr(derived, f"l{i}_inv_kv_scales") for i in range(self.num_layers) + } + self.moe = { + i: _MoeScalars( + **{ + f.name: getattr(derived, f"l{i}_{f.name}") + for f in dataclasses.fields(_MoeScalars) + } + ) + for i in sorted(self.moe_layers) + } + self._build_embed_owners() + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self._built_for = identity + + def cache_derived_state(self) -> None: + """The staged loaders' entry to the same build (``post_load_weights``).""" + self.post_load_weights() + + def _derive_moe(self, i: int) -> _MoeScalars: + """The routed experts' kernel scalars from the checkpoint's per-tensor scales (trtllm-gen NVFP4 layout). + + Global input scales: the reciprocal of the largest per-expert input scale (gate and up must agree). Per expert: + ``alpha = 1 / (input_scale * (1 / weight_scale_2))`` for the gate (fc31_alpha), the up half (its GEMM2-input + rescale ``fc31_scale_c = fc2_input_scale * up_alpha``) and the down projection; SwiGLU-OAI's activation + constants per expert, with beta and the clamp in the pre-dequant GEMM output domain (divided by fc31_alpha). + """ + w = self.w + isc1, isc3, isc2 = w[f"l{i}_e_isc1"].data, w[f"l{i}_e_isc3"].data, w[f"l{i}_e_isc2"].data + if not torch.allclose(isc1, isc3): + raise ValueError(f"layer {i}: the routed experts' w1 and w3 input scales differ") + fc31_input_scale = isc1.max().reciprocal() + fc2_input_scale = isc2.max().reciprocal() + ws2_1, ws2_3, ws2_2 = ( + w[f"l{i}_e_ws2_1"].data, + w[f"l{i}_e_ws2_3"].data, + w[f"l{i}_e_ws2_2"].data, + ) + fc31_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_1)) + up_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_3)) + fc2_alpha = 1.0 / (fc2_input_scale * (1.0 / ws2_2)) + fc31_scale_c = fc2_input_scale * up_alpha + n = self.num_experts + device = fc31_alpha.device + act_alpha = torch.full((n,), self.swiglu_alpha, dtype=torch.float32, device=device) + act_beta = torch.full((n,), self.swiglu_beta, dtype=torch.float32, device=device) + act_clamp = torch.full((n,), self.swiglu_limit, dtype=torch.float32, device=device) + act_beta.div_(fc31_alpha) + act_clamp.div_(fc31_alpha) + return _MoeScalars( + fc31_input_scale=fc31_input_scale, + fc2_input_scale=fc2_input_scale, + fc31_alpha=fc31_alpha, + fc2_alpha=fc2_alpha, + fc31_scale_c=fc31_scale_c, + act_alpha=act_alpha, + act_beta=act_beta, + act_clamp=act_clamp, + ) + + def _build_embed_owners(self) -> None: + """The fused embedding + first RMSNorm's caller-owned buffers, one per (rows, id dtype) of a decode step.""" + table, weight = self.w["embed"], self.w["l0_in_norm"] + self.embed_owners = { + (rows, dtype): DecodeEmbedding(table, weight, self.rms_eps, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + + def projections(self) -> List[Tuple[str, str, Optional[str]]]: + """Every MXFP8 projection in layer order: (weight key, scale key, the boundary its GEMV pushes into).""" + out = [] + for i in range(self.num_layers): + out.append((f"l{i}_qkv_w", f"l{i}_qkv_s", None)) + out.append((f"l{i}_o_w", f"l{i}_o_s", "attn")) + if i in self.moe_layers: + out.append((f"l{i}_sh_gu_w", f"l{i}_sh_gu_s", None)) + out.append((f"l{i}_sh_dn_w", f"l{i}_sh_dn_s", "moe")) + else: + out.append((f"l{i}_mlp_gu_w", f"l{i}_mlp_gu_s", None)) + out.append((f"l{i}_mlp_dn_w", f"l{i}_mlp_dn_s", "dense")) + return out + + def decode_query_lens(self) -> List[int]: + """The query rows per request a decode step can have: 1, and 1 + each draft length speculation runs.""" + lens = {1} + spec = getattr(self.model_config, "spec_config", None) + if spec is not None: + for drafts in ( + getattr(spec, "max_draft_len", None), + getattr(spec, "max_total_draft_tokens", None), + *(getattr(spec, "draft_len_schedule", None) or {}).values(), + ): + if drafts: + lens.add(1 + int(drafts)) + return sorted(lens) + + def prebuild_kernels(self) -> Dict[str, object]: + """Compile (and load) every kernel a decode forward dispatches to that would otherwise compile on its first + call: each projection's decode GEMVs (every row count, activation and boundary push), the MoE engines of + ``MOE_ENGINE``'s token counts, the index scorers at the decode query lengths (this target's and the in-tree + one wider steps take), the 384-thread long-row select, the dense layers' full-context decode plans, and the + built-in sparse decode a mixed step's generation rows run. Runs on the current (main) stream before the + executor's warmup, in the same order on every rank, then synchronizes; after it, GEMV compiles are frozen and + the forward only looks kernels up (a missing one is a miss). Returns the seconds each part took and the keys. + """ + times: Dict[str, object] = {} + # The 9..16-row generation router: decided once, from the device, never inside a CUDA-graph capture. + # The built-in gate's conditions, mirrored exactly: 128 experts, a hidden size in 1024s, 16-CTA clusters. + cfg = self.model_config.pretrained_config + self.router_tc_cluster16 = ( + int(cfg.num_local_experts) == 128 + and int(cfg.hidden_size) % 1024 == 0 + and bool(torch.ops.trtllm.minimax_m3_router_tc_cluster16(int(cfg.hidden_size))) + ) + logger.info( + f"[M3 prebuild] rank {mpi_rank()}: TF32 tensor-core router for {ROUTER_TC_GEN_MIN_ROWS}.." + f"{ROUTER_TC_GEN_MAX_ROWS}-row generation steps: " + + ( + "engaged (16-CTA clusters)" + if self.router_tc_cluster16 + else "not used (no 16-CTA clusters)" + ) + ) + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + from .kernels import m3_select_long_dsl as _select_long_dsl_kernel + from .kernels import m3_tokmap_guard + + device = self.w["embed"].device + # The one-request token-map guard of the stock generation preprocess (``PREPROCESS_TOKMAP_FIX``): every path + # that can run it (the drafter's full-context decode or its fallback, any trtllm-gen fallback), so it is + # installed whatever else is enabled. + m3_tokmap_guard.install( + torch.device("cuda", torch.cuda.current_device()), + DECODE_MAX_TOKENS, + lambda rows: taken(PREPROCESS_TOKMAP_PATH, rows, "one-request token map"), + lambda rows, reason: missed(PREPROCESS_TOKMAP_PATH, rows, reason), + ) + mx_gemv.freeze(False) + use_pdl = get_env_enable_pdl() + world = int(self.mapping.tp_size) + start = time.perf_counter() + for wk, sk, push in self.projections(): + weight, scale = self.w[wk], self.w[sk] + mx_gemv.ensure_compiled(weight, scale, use_pdl) + n, k = weight.shape + for rows, x_bf16, ar in _gemv_variants(n, k, push, world): + mx_gemv.precompile(weight, scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + torch.cuda.synchronize() + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for i in sorted(self.moe_layers): + self._prebuild_engines(i) + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = self._prebuild_index_scorer(self.decode_query_lens(), device) + torch.cuda.synchronize() + times["scorer_s"] = time.perf_counter() - start + start = time.perf_counter() + times["select_keys"] = [] + if long_select_dsl_keys(): + # The 384-thread long-row select (``SELECT_LONG_DSL``): one kernel per launch key. + times["select_keys"] = _select_long_dsl_kernel.prebuild(long_select_dsl_keys(), device) + torch.cuda.synchronize() + times["select_s"] = time.perf_counter() - start + start = time.perf_counter() + times["fullctx_plans"] = 0 + if FULLCTX_DECODE.keys: + # The dense layers' full-context decode: one compiled PrimTS plan per covered (rows, requests). + times["fullctx_plans"] = _fullctx_decode_kernel.prebuild( + device, + {key: FULLCTX_DECODE_RECIPES[key] for key in FULLCTX_DECODE.keys}, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), + use_pdl, + ) + torch.cuda.synchronize() + times["fullctx_s"] = time.perf_counter() - start + start = time.perf_counter() + times["builtin_decode_rows"] = _prebuild_builtin_sparse_decode( + self.decode_query_lens(), self.heads, self.kv_heads, device + ) + torch.cuda.synchronize() + times["builtin_decode_s"] = time.perf_counter() - start + mx_gemv.freeze(True) + logger.info( + "[M3 prebuild] decode kernels built before warmup: GEMVs %.1f s, MoE engines %.1f s, index scorer " + "%.1f s (%d variants), long-row select %.1f s (%d variants), full-context decode %.1f s (%d plans), " + "built-in sparse decode %.1f s (%d row counts)" + % ( + times["gemv_s"], + times["moe_s"], + times["scorer_s"], + len(times["scorer_keys"]), + times["select_s"], + len(times["select_keys"]), + times["fullctx_s"], + times["fullctx_plans"], + times["builtin_decode_s"], + len(times["builtin_decode_rows"]), + ) + ) + return times + + # ------------------------------------------------------------------------------------------ the MoE engines + + def _expert_operands(self, i: int): + """The routed experts' operands the engines read (trtllm-gen NVFP4 SwiGLU-OAI layout), and the input scale.""" + w, m = self.w, self.moe[i] + tensors = [ + w[f"l{i}_e_w31"].data, + w[f"l{i}_e_s31"].data, + w[f"l{i}_e_w2"].data, + w[f"l{i}_e_s2"].data, + m.fc31_alpha, + m.fc31_scale_c, + m.act_alpha, + m.act_beta, + m.act_clamp, + m.fc2_alpha, + ] + return tensors, m.fc31_input_scale + + def _shared_weights(self, i: int): + w = self.w + return ( + w[f"l{i}_sh_gu_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_gu_s"].data, + w[f"l{i}_sh_dn_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_dn_s"].data, + ) + + def engine(self, i: int, num_tokens: int, shared: bool) -> Optional[object]: + """Layer ``i``'s prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in (and the shared + expert with ``shared``). None after a miss: the build failed or was never made (a forward never builds one).""" + return self._lookup( + i, self.moe_engines[i], (num_tokens, True, shared), MOE_ENGINE.path, num_tokens + ) + + def _lookup(self, i: int, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {i}'s engine" + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, what) + else: + missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + missed(path, num_tokens, "its build failed (logged)") + return engine + + def _prebuild_engines(self, i: int) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of layer ``i`` can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in. Whether an engine is kept is agreed + across the TP group (one host collective per engine): a build that fails on one rank drops that engine on + every rank, so the ranks never mix an engine (and its boundary push) with the generic MoE.""" + engines = self.moe_engines.setdefault(i, {}) + for num_tokens in sorted(MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in engines: + engines[key] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build( + i, MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + f"[M3 dispatch] {MOE_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + + def _try_build(self, i: int, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(i, num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {i}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, i: int, num_tokens: int, shared: bool) -> object: + """The 1..8-token engine (router, routed experts and the shared expert in one launch, at the TP2 width 1536). + Rubin TP2 ablation chain (57 layers behind a 58.4 MB PDL producer, row = (P + MoE) - P, both orders): 31.3 / + 50.7 us/layer at 4 / 8 tokens against the built-in MoE's 42.4 / 55.4 (FlashInfer MXFP8 shared expert); bit for + bit the built-in block's shared.add_(routed). When the engine pushes its sum into the boundary all-reduce, + every CTA releases the PDL dependents right after its grid dependency wait instead of after the push + (``MOE_EARLY_AR_RELEASE``), so the all-reduce (which polls the pushed tagged words, not the grid) becomes + resident on the SMs the engine frees: scheduling only, the same bits (Rubin TP2 greedy early column, same-job + A,B + B,A: c1 -15 us).""" + tensors, scale = self._expert_operands(i) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + kwargs["shared_weights"] = self._shared_weights(i) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = M3MoeState( + device, + num_tokens=num_tokens, + num_ctas=sms - MOE_ENGINE_RESERVED_SMS, + shared=shared, + early_release=MOE_EARLY_AR_RELEASE, + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + # ------------------------------------------------------------------------------------------ the index scorers + + def _index_geometries(self) -> list: + """Every sparse layer's (index dtype, index heads per rank, index dim), deduplicated.""" + geoms = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + geom = (dtype, int(config.num_index_heads), int(config.sparse_index_dim)) + if geom not in geoms: + geoms.append(geom) + return geoms + + def _scorer_variants(self, query_lens) -> list: + """The scorer keys a decode step can take: every sparse layer's index geometry at each decode query length, + with each covered request count's schedule (``index_score_schedule``), where the kernel serves it.""" + runner = _index_score_kernel.IndexScoreRunner + keys = [] + for dtype, heads, dim in self._index_geometries(): + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=heads, + head_dim=dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for requests in sorted(INDEX_SCORE.keys): + split_k, early = index_score_schedule(requests) + key = (dtype, heads, int(qlen), dim, int(split_k), bool(early)) + if key not in keys: + keys.append(key) + return keys + + def _in_tree_scorer_variants(self, query_lens) -> list: + """The in-tree CuTe DSL scorer's keys (``CuteDSLMiniMaxM3IndexDecodeScoreRunner``: index dtype, index heads, + decode query length, index dim, its 256-CTA split) a decode step outside ``INDEX_SCORE`` can take: every query + length at which a step of more than ``INDEX_SCORE_MAX_REQUESTS`` requests fits in ``DECODE_MAX_TOKENS`` (e.g. + 3..8 requests decoding one row each), where that kernel serves the geometry.""" + in_tree = _in_tree_scorer() + widest = DECODE_MAX_TOKENS // (INDEX_SCORE_MAX_REQUESTS + 1) + keys = [] + for dtype, heads, dim in self._index_geometries(): + for qlen in sorted(set(query_lens)): + if qlen > widest or not in_tree.is_supported( + q_dtype=dtype, + num_heads=heads, + head_dim=dim, + page_size=in_tree.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + key = (dtype, heads, qlen, dim, in_tree.SPLIT_K) + if key not in keys: + keys.append(key) + return keys + + def _prebuild_index_scorer(self, query_lens, device) -> list: + """Compile every key of ``_scorer_variants`` and ``_in_tree_scorer_variants`` and launch each once on + real-shaped dummy operands (loading its module), on the current stream, so no decode step compiles a scorer. + The kernels' batch, token, page and block counts are symbolic: their compiles need no step metadata. Returns + the keys.""" + runner = _index_score_kernel.IndexScoreRunner + keys = self._scorer_variants(query_lens) + for key in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + q, k, bt, score, seq = _dummy_scorer_operands( + dtype, heads, qlen, dim, runner.SUPPORTED_PAGE_SIZE, device + ) + runner.kernel_cache[key](q, k, bt, score, seq) + in_tree = _in_tree_scorer() + in_tree_keys = self._in_tree_scorer_variants(query_lens) + for key in in_tree_keys: + dtype, heads, qlen, dim, split_k = key + in_tree._compile(dtype, heads, qlen, dim, split_k) + q, k, bt, score, seq = _dummy_scorer_operands( + dtype, heads, qlen, dim, in_tree.SUPPORTED_PAGE_SIZE, device + ) + in_tree.kernel_cache[key](q, k, bt, score, seq) + return keys + in_tree_keys + + # --------------------------------------------------------------------------------------------- the dispatch + + def _select_target(self, attn_metadata: AttentionMetadata): + """Which target runs this step: one host-side integer, fixed per CUDA-graph capture. A step with context rows + (a mixed step included) is prefill's; every pure generation step is decode's, whatever its width.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + capture: Callable[..., None] = _plain_capture, + capture_target: Optional[Callable[..., Optional[tuple]]] = None, + **kwargs, + ) -> torch.Tensor: + """``capture``: the Eagle3 hidden-state capture of a layer (the worker's, with its capture norm), and + ``capture_target``: where a decode boundary all-reduce writes a capture layer's capture norm itself + (``CAPTURE_NORM_IN_BOUNDARY``; the worker's, None without one).""" + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if kwargs.get("lora_params") is not None: + raise ValueError("this MiniMax-M3 target serves no LoRA adapters") + return self._select_target(attn_metadata).forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The targets +# -------------------------------------------------------------------------------------------------------------------- + + +#: The cores the breakable prefill graphs' eager attention boundary looks up by key: an ``eager_on_graph`` callable +#: replays with weak references to its arguments, so they may be tensors or scalars only. +_BCG_CORES: Dict[str, "MiniMaxM3Tp2ModelingV2"] = {} + + +def _msa_attention_live(core, i: int, q, k, v, idx_q, md, output) -> None: + """The MSA backend over a step's live rows (``md.num_tokens``; a padded step's pad rows of ``output`` are left as + allocated): the KV write unless a producer wrote the caches, the sparse layers' top-k block selection, the FMHA + over the context rows and any generation rows beside them.""" + live = int(md.num_tokens) + sparse = i in core.sparse_layers + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k is not None) + m3_msa_attention( + core.attn[i], + q[:live], + k[:live] if k is not None else None, + v[:live] if v is not None else None, + idx_q[:live] if idx_q is not None else None, + md, + output[:live], + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + + +def _prefill_attention_boundary( + packed, position_ids, q, k, v, idx_q, output, core_key: str, i: int +) -> None: + """The attention boundary of a prefill layer under a breakable prefill CUDA graph, run eagerly at every replay (as + the built-in model's MiniMax-M3 attention boundary is): everything that depends on the step's live row count. A + sparse layer passes its ``packed`` five-way projection: the horizontal producer runs here over the live rows + (K / V / index-K into the paged caches, Q / index-Q out), then the MSA backend; a dense layer passes q / k / v.""" + core = _BCG_CORES[core_key] + md = get_model_extra_attrs()["attention_metadata"]() + live = int(md.num_tokens) + if packed is not None: + horizontal = _sparse_qkv_producer( + core, + i, + packed[:live], + position_ids[..., :live] if position_ids is not None else None, + md.kv_cache_manager, + md.msa_idx_k_cache, + md.msa_out_cache_loc, + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + k = v = None + _msa_attention_live(core, i, q, k, v, idx_q, md, output) + + +_prefill_attention_eager = eager_on_graph(_prefill_attention_boundary) + +#: The drafters the breakable prefill graphs' eager draft attention boundary looks up by key (``_BCG_CORES``'s rule). +_BCG_DRAFTS: Dict[str, "M3Eagle3Draft"] = {} + + +def _draft_attention_boundary(qkv, output, draft_key: str) -> None: + """The drafter layer's attention under a breakable prefill CUDA graph, run eagerly at every replay (as the in-tree + drafter's attention boundary is): the backend's RoPE + KV append + causal attention over the step's live rows, + with the step's metadata read at the replay; ``output``'s pad rows are left as allocated.""" + draft = _BCG_DRAFTS[draft_key] + md = get_model_extra_attrs()["attention_metadata"]() + out = m3_draft_attention(draft.attn, qkv, md) + output[: out.shape[0]].copy_(out) + + +_draft_attention_eager = eager_on_graph(_draft_attention_boundary) + + +class PrefillTarget: + """A step with context rows: prefill chunks, and the mixed steps that carry generation rows beside them. Every + layer runs in its generic form; the attention core is the MSA backend's for every row. + + Under a breakable (piecewise) prefill CUDA graph the projections, norms, MoE and boundaries are captured per padded + row bucket and the attention boundary runs eagerly at every replay (``_prefill_attention_eager``: the live-row + horizontal producer and the MSA backend), as the built-in model's MiniMax-M3 attention does.""" + + def __init__(self, core: MiniMaxM3Tp2ModelingV2) -> None: + self.core = core + self._bcg_key = f"m3_tp2_core_{id(core)}" + _BCG_CORES[self._bcg_key] = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + Counts and per-step host scalars are read here; device tensors are the metadata's persistent buffers, held + by reference; ``msa_idx_k_cache`` is the metadata's per-layer index-K cache getter. ``metadata`` is the + prepared object itself: the MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_tokens=md.num_tokens, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + """``capture_target`` is the decode boundaries' (unused: no prefill boundary writes a capture norm).""" + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + _record("prefill", int(step["num_tokens"]), int(step["num_contexts"])) + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + hidden, residual = self._layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + return hidden + + def _swiglu_mlp( + self, h: torch.Tensor, prefix: str, low_m_gemm: bool, dense: bool = False + ) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI (the dense MLP's on this target's + C++ op, ``_dense_swiglu``), down, no reduction.""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + out, _ = _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm) + return out + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: + the MoE scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: on ``ROUTER_TC_GEN``'s rows the TF32 tensor-core router with the + experts' NVFP4 activation fused in (as the agg canon routes a 9..48-row step with context rows; one runner + call), else the fp32 router logits and the experts quantizing their bf16 input (chunked).""" + core = self.core + fused = _router_tc_fp4(core, i, h) + if fused is not None: + logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, logits, True) + return self._routed_chunks(i, h, _router_logits(core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool): + """The MoE block of layer ``i`` without its reduction. Returns ``(out, routed)``: with ``return_parts`` the + shared and routed parts (routed None when the engine folded the shared expert in), else their sum in + ``out``.""" + core = self.core + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, False) + taken(MOE_ENGINE.path, rows) + folded = m3_moe_engine(h.contiguous(), engine, None, None) + return folded, None + if not return_parts and _triton_finalize_rows(core, h): + routed_parts, shared_out = maybe_execute_in_parallel( + lambda: _routed_experts(core, i, h, None, _router_logits(core, i, h), False), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + return _triton_finalize(rows, routed_parts, shared_out), None + routed_out, shared_out = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the + # same collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + # Row ranges at this dispatch: 9..48 rows take the decode target's generic MC finalize boundary (as the agg + # canon's generic layer does on every step, context rows included: the MoE unfinalized, its finalize + + # shared expert in the multicast all-reduce + residual + next norm); the CTX stack starts at 4096 rows. + folded = core._decode._generic_mc_finalize(i, hidden, residual, _next_norm(core, i)) + if folded is not None: + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, folded[1], None) + return folded + out, routed = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step): + """Attention of layer ``i`` up to its o_proj (no reduction). Returns (o_proj output, pushed): a prefill step + arms no boundary push, so pushed is always False.""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v for an FP8 cache, on this target's op; else in + # place in bf16), the MSA backend (it writes the layer's cache), o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = ( + _dense_qk_norm_rope_fp8(core, i, qkv, position_ids) + if core.main_kv_is_fp8 + else _qk_norm_rope(core, i, qkv, position_ids, False) + ) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + if not core.fused_projection: + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target runs a breakable prefill CUDA graph only with the five-way projection " + "(fuse_qkv_index_projection)" + ) + return self._unfused_sparse_attention(i, position_ids, hidden, step) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + if is_in_breakable_cuda_graph(): + output = packed.new_empty( + (packed.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16 + ) + _prefill_attention_eager( + packed, position_ids, None, None, None, None, output, self._bcg_key, i + ) + return _linear(core, output, f"l{i}_o_w", f"l{i}_o_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection + and its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and + the FP8 indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the + o_proj. A layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): the MSA + backend (``m3_msa_attention``: the KV write unless a producer wrote the caches, the sparse layers' top-k block + selection, the FMHA over the context rows and any generation rows beside them).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + if is_in_breakable_cuda_graph(): + # The live row count is a host value of the step: read at every replay inside the eager boundary. + _prefill_attention_eager(None, None, q, k, v, idx_q, output, self._bcg_key, i) + else: + _msa_attention_live(core, i, q, k, v, idx_q, step["metadata"], output) + return output + + +class DecodeTarget: + """A pure generation step (no context rows), of any width. + + Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids (up to 8 requests verifying 1 + 3 rows), per layer: + attention (o_proj all-reduce deferred, its rows pushed into the boundary all-reduce up to 16 rows) -> attention + boundary (all-reduce + residual + post-attention RMSNorm) -> MoE (the 1..8-row engine pushing its sum into the + boundary all-reduce; at 9..32 rows the unfinalized MoE with the finalize in the boundary all-reduce) or dense MLP + (its down projection pushing into the post-MLP boundary) -> MoE / dense boundary (all-reduce + shared/routed add + + residual + the next layer's RMSNorm, emitting the next QKV's MXFP8 activation). An Eagle3 capture layer runs the + same decode layer and captures its feed-forward boundary's reduced residual sum (the boundary writing the capture + norm itself when the worker gives it a target). The first layer runs the generic layer when the fused embedding + does not take the step. A wider step (or one carrying embeddings) runs every layer in its generic form, with this + target's attention (the dense layers' full-context decode, the decode selection and sparse decode where the + producer wrote the caches).""" + + def __init__(self, core: MiniMaxM3Tp2ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # Set by a feed-forward boundary that wrote the layer's Eagle3 capture (see _layer's ``capture``). + self._capture_written = False + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + States the two context counts instead of reading them back: this target runs only when ``num_contexts`` is + zero, so reading them would ask the metadata to confirm what the routing already decided. + + Graph classes: counts, the generation span, the identity-selection flag, the worst-case score extent and the + cache length bound are host values fixed per capture (prepare() stages them before the forward); the block + table, lengths, cache slots, identity table and valid-block counts are the metadata's persistent buffers, held + by reference and refreshed in place every step; ``msa_idx_k_cache`` / ``msa_proxy_max_score_view`` are the + metadata's getters of the per-layer index-K cache and of the score scratch view. ``metadata`` is the prepared + object itself: the MSA backend's step (``m3_msa_attention``) and the full-context decode's K/V write + (``write_msa_phase_kv``) take it as their argument. + """ + return dict( + num_contexts=0, + num_ctx_tokens=0, + num_tokens=md.num_tokens, + num_seqs=md.num_seqs, + max_seq_len=md.max_seq_len, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + msa_decode_span=md.msa_decode_span, + msa_block_table=md.msa_block_table, + msa_seq_lens_cuda=md.msa_seq_lens_cuda, + msa_short_ok=md._msa_short_ok, + msa_short_table=md.msa_short_table, + msa_proxy_max_score_view=md.msa_proxy_max_score_view, + msa_worst_case_max_k_tiles=md.msa_worst_case_max_k_tiles, + msa_n_valid_blocks=md.msa_n_valid_blocks, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + # The decode kernels' rows (a per-capture constant): up to DECODE_MAX_TOKENS tokens of 1-D input ids. + kernels = ( + inputs_embeds is None and input_ids.ndim == 1 and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + _record("decode", num_tokens, plan.name) + hidden, residual = self._embed(input_ids, num_tokens) + else: + # Breakable (piecewise) CUDA graphs capture steps with context rows only (prefill_cuda_graph_backend): + # every pure generation step, the 33..192-row ones on this generic branch included, runs eagerly or on + # the ordinary decode CUDA graphs, so this guard never fires in a supported deployment. + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + _record("decode", int(step["num_tokens"]), "generic") + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + captured = spec_metadata is not None and spec_metadata.is_layer_capture(i) + if kernels and residual is not None and captured: + # The capture layer on the decode layer (its boundaries push); the capture reads the MoE / dense + # boundary's reduced residual sum, as the generic layer's fused boundary capture does. Given a target + # (``CAPTURE_NORM_IN_BOUNDARY``), the boundary all-reduce writes the capture norm itself. + target = ( + capture_target(spec_metadata, i, num_tokens) + if capture_target is not None + else None + ) + self._capture_written = False + hidden, residual = self._layer( + i, plan, position_ids, hidden, residual, step, capture=target + ) + taken("capture_layer_fast", num_tokens, f"layer {i}") + if target is not None and self._capture_written: + taken("capture_norm_in_boundary", num_tokens, f"layer {i}") + capture(spec_metadata, i, residual, None, written=True) + else: + capture(spec_metadata, i, residual, None) + elif not kernels or residual is None or captured: + hidden, residual = self._generic_layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + else: + hidden, residual = self._layer(i, plan, position_ids, hidden, residual, step) + return hidden + + def _embed(self, input_ids: torch.Tensor, num_tokens: int): + """(hidden, residual) of the fused embedding + first RMSNorm; (embedding, None) when it does not run, so the + first layer runs its own input norm. SM107: 1.55..1.65 us against 3.4..6.0 us for the stock embedding + + RMSNorm at 1..8 rows (the gather bitwise, the norm one ulp off on ~5e-6 of the elements: summation order).""" + core = self.core + if not EMBED_FIRSTNORM.covers(num_tokens): + uncovered(EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens") + return m3_embedding(input_ids, core.w["embed"]), None + owner = core.embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + missed( + EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 after the load)" + if owner is None + else "the input ids are not contiguous", + ) + return m3_embedding(input_ids, core.w["embed"]), None + _record("m3_embedding_firstnorm", num_tokens, input_ids.dtype) + taken(EMBED_FIRSTNORM.path, num_tokens, str(input_ids.dtype)) + return m3_embedding_firstnorm( + input_ids, owner.table, owner.weight, owner.eps, owner.raw, owner.output, owner.use_pdl + ) + + def _bucket_flags(self, num_tokens: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(self.core.mapping, self.core.hidden) + workspace = None if workspace is None else m3_mc_allreduce._bucket(workspace, num_tokens) + flags = None if workspace is None else workspace.flags_for(8) + return None if flags is None else (workspace, flags) + + @staticmethod + def _no_bucket(path: str, num_tokens: int) -> None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, num_tokens: int): + """The push tuple into the attention boundary's all-reduce (the o_proj's; also the dense down projection's into + the post-MLP boundary, the same exchange), or None (after a miss). + + None comes only from the token count or from workspace state created collectively in the same eager call on + every rank, so all ranks agree. (A lone non-pushing rank would still be safe: its all-reduce pushes the same + bf16 rows into the same set A.)""" + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(ATTN_PUSH.path, num_tokens) + return None + workspace, flags = found + return ( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + AR_PUSH_MODE, + ) + + def _moe_push_targets(self, num_tokens: int): + """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None + (after a miss). + + This decision must match across ranks: prepushed 3 reads set B of every rank, which a non-pushing rank never + fills. None comes only from the token count or from state created in the same eager call on every rank + (workspace, flags, the descriptor cache), and the engine it is armed for is agreed at build time + (``_prebuild_engines``), so it does.""" + hidden = self.core.hidden + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(MOE_PUSH.path, num_tokens) + return None + workspace, flags = found + key = (workspace.mc.data_ptr(), flags.data_ptr()) + desc = self._moe_push_descs.get(key) + if desc is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor") + return None + desc = self._moe_push_descs[key] = ar_push_desc( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + flags.device, + ) + set_b = workspace.mc.data_ptr() + 2 * 4 * workspace.max_tokens * workspace.world_size * ( + hidden // 2 + ) + return desc, ( + set_b, + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + 2, + ) + + def _boundary(self, use_mc: bool, x, residual, norm_w, x2=None, quant=False, capture=None): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce. + ``capture``: the Eagle3 capture the multicast kernel also writes (``_capture_written`` records it).""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant, capture=capture) + if fused is not None: + if capture is not None: + self._capture_written = True + taken( + MC_ALLREDUCE.path, + x.shape[0], + ("x + x2" if x2 is not None else "x") + (", MXFP8 copy" if quant else ""), + ) + return fused + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + missed( + MC_ALLREDUCE.path, + x.shape[0], + m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + ) + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _prepushed( + self, i: int, x, residual, norm_w, quant, prepushed, source, x2=None, capture=None + ): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed, capture=capture + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + if capture is not None: + self._capture_written = True + taken(MC_ALLREDUCE.path, x.shape[0], f"prepushed {prepushed}, {source}") + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step, capture=None): + """Layer ``i`` on the decode kernels. ``capture``: (out, fc_norm weight, eps) for the feed-forward boundary to + write the layer's Eagle3 capture norm (``CAPTURE_NORM_IN_BOUNDARY``); a boundary that writes it sets + ``_capture_written``.""" + core = self.core + num_tokens = hidden.shape[0] + # Padded rows: the boundary all-reduce polls and re-arms rows t < M = residual.shape[0], CUDA-graph padding + # rows included. Every producer pushes all rows of the same tensor: the o_proj / shared down / dense down GEMV + # t < m = x.shape[0] (its row count comes from the host shape, not a live count), the MoE engine every row of + # the num_tokens it was built for (= num_tokens here). The all-reduce checks x.sizes() == residual.sizes(), so + # no padded row is left with a stale flag. + # The GEMV dispatch (and with it every push into a boundary all-reduce) is a decode-graph one. + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + uncovered( + "decode_graph", + "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + ) + push = plan.attn_mc_ar and decode_graph and ATTN_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if push else None + hidden, attn_pushed = self._attention(i, position_ids, hidden, step, armed) + post_norm = core.w[f"l{i}_post_norm"] + if armed is not None and attn_pushed: + hidden, residual = self._prepushed( + i, hidden, residual, post_norm, False, AR_PUSH_MODE, "o_proj push" + ) + else: + if armed is not None: + missed( + ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden, residual = self._boundary(plan.attn_mc_ar, hidden, residual, post_norm) + + norm_w = _next_norm(core, i) + if i not in core.moe_layers: + # A dense layer: its MLP's down projection pushes into the post-MLP boundary's all-reduce, as the o_proj + # does into the attention one (the same exchange; the gate_up projection never pushes). SM107 TP2: the + # two non-capture dense layers' post-MLP hops end 5.5-5.8 us after the down GEMV unpushed vs 2.4-2.5 us + # pushed; bitwise equal to the unpushed boundary (the same bf16 words). + dense_push = plan.attn_mc_ar and decode_graph and DENSE_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if dense_push else None + out, pushed = self._swiglu_mlp(hidden, f"l{i}_mlp", True, armed, dense=True) + if armed is not None and pushed: + return self._prepushed( + i, + out, + residual, + norm_w, + False, + AR_PUSH_MODE, + "dense down push", + capture=capture, + ) + if armed is not None: + missed( + DENSE_PUSH.path, + num_tokens, + "the dense MLP's down projection did not push its rows into the armed boundary all-reduce", + ) + return self._boundary(plan.attn_mc_ar, out, residual, norm_w, capture=capture) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + moe_push = None + if moe_path(num_tokens) == MOE_ENGINE.path: + # The engine runs inside the MoE; looked up (or found missing) here first so a push is armed only for an + # engine that will take it. + engine = core.engine(i, num_tokens, shared=True) + if ( + engine is not None + and plan.moe_mc_ar + and decode_graph + and MOE_PUSH.covers(num_tokens) + ): + moe_push = self._moe_push_targets(num_tokens) + elif moe_path(num_tokens) == MC_FINALIZE.path and plan.moe_mc_ar: + if self._bucket_flags(num_tokens) is None: + self._no_bucket(MC_FINALIZE.path, num_tokens) + else: + return self._mc_finalize(i, hidden, residual, norm_w, quant, capture) + shared, routed, engine_pushed, shared_pushed = self._moe( + i, hidden, return_parts=True, push=moe_push + ) + if moe_push is not None: + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + if engine_pushed: + return self._prepushed( + i, shared, residual, norm_w, quant, 2, "MoE engine push", capture=capture + ) + else: + if engine_pushed and shared_pushed: + return self._prepushed( + i, + shared, + residual, + norm_w, + quant, + 3, + "MoE engine + shared push", + x2=routed, + capture=capture, + ) + if engine_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + missed( + MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary( + plan.moe_mc_ar, shared, residual, norm_w, x2=routed, quant=quant, capture=capture + ) + + def _mc_finalize(self, i: int, h, residual, norm_w, quant: bool, capture=None, what: str = ""): + """The MoE of layer ``i`` unfinalized (its shared expert beside it) and the finalize + all-reduce + residual + + next RMSNorm in one launch (``m3_mc_finalize_allreduce_norm``): returns (normed, updated residual). The + caller checked the rows and the multicast bucket (decided from shapes alone, so the same on every TP rank).""" + core = self.core + rows = h.shape[0] + routed, (shared, _) = maybe_execute_in_parallel( + lambda: _decode_routed_unfinalized(core, i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + fc2, scales, perm = routed + _record("m3_mc_finalize_allreduce_norm", rows, quant, capture is not None) + fused = m3_mc_allreduce.m3_mc_finalize_allreduce_norm( + fc2, + perm.view(rows, -1).to(torch.int32), + scales.reshape(rows, -1), + shared, + residual, + norm_w, + core.rms_eps, + m3_mc_allreduce.get_workspace(core.mapping, core.hidden), + use_pdl=get_env_enable_pdl(), + quant=quant, + capture=capture, + ) + if fused is None: + raise RuntimeError( + f"layer {i}: the MC finalize all-reduce rejected the unfinalized MoE outputs" + ) + if capture is not None: + self._capture_written = True + taken(MC_FINALIZE.path, rows, what or ("MXFP8 copy" if quant else "")) + return fused + + def _generic_mc_finalize(self, i: int, h, residual, norm_w): + """Past ``DECODE_MAX_TOKENS`` rows, up to ``GENERIC_MC_FINALIZE_MAX_TOKENS``: the generic layer's MoE boundary + as ``_mc_finalize``; None (from shapes alone, the same on every rank) keeps the generic boundary.""" + rows = h.shape[0] if h.dim() == 2 else 0 + cap = min(GENERIC_MC_FINALIZE_MAX_TOKENS, m3_mc_allreduce.FINALIZE_MAX_TOKENS) + if ( + not (DECODE_MAX_TOKENS < rows <= cap) + or h.dtype != torch.bfloat16 + or torch.compiler.is_compiling() + ): + return None + if self._bucket_flags(rows) is None: + return None + return self._mc_finalize(i, h, residual, norm_w, False, what="generic") + + def _swiglu_mlp( + self, + h: torch.Tensor, + prefix: str, + low_m_gemm: bool, + push: Optional[tuple] = None, + dense: bool = False, + ): + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI (the dense MLP's on this target's + C++ op, ``_dense_swiglu``), down, no reduction. Returns (output, whether the down projection's GEMV pushed into + ``push``).""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm, push) + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: the MoE + scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: on ``ROUTER_TC_GEN``'s rows the TF32 tensor-core router with the + experts' NVFP4 activation fused in (one runner call: the rows are within ``moe_max_num_tokens``), else the + fp32 router and the experts quantizing their bf16 input (chunked).""" + core = self.core + fused = _router_tc_fp4(core, i, h) + if fused is not None: + logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, logits, True) + return self._routed_chunks(i, h, _decode_router_logits(core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool, push=None): + """The MoE block of layer ``i`` without its reduction. + + ``push`` = (the engine's ``arp`` descriptor, the shared down projection's ``ar`` tuple) the decode layer armed + for the MoE boundary's exchange sets, or None. Returns ``(out, routed, engine_pushed, shared_pushed)``: with + ``return_parts`` the shared and routed parts (routed None when the engine folded the shared expert in), else + their sum in ``out``. + """ + core = self.core + engine_arp, shared_ar = push if push is not None else (None, None) + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, engine_arp is not None) + taken(MOE_ENGINE.path, rows) + folded = m3_moe_engine(h.contiguous(), engine, None, engine_arp) + return folded, None, engine_arp is not None, False + if not return_parts and push is None and _triton_finalize_rows(core, h): + routed_parts, (shared_out, _) = maybe_execute_in_parallel( + lambda: _decode_routed_unfinalized(core, i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + return _triton_finalize(rows, routed_parts, shared_out), None, False, False + routed_out, (shared_out, shared_pushed) = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False, shared_ar), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the same + # collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out, False, shared_pushed + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None, False, shared_pushed + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + if x.shape[0] > DECODE_MAX_TOKENS: + taken( + MC_ALLREDUCE.path, + x.shape[0], + "generic " + ("x + x2" if x2 is not None else "x"), + ) + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _generic_layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + hidden, residual = self._generic_boundary(hidden, None, residual, core.w[f"l{i}_post_norm"]) + if i in core.moe_layers: + folded = self._generic_mc_finalize(i, hidden, residual, _next_norm(core, i)) + if folded is not None: + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, folded[1], None) + return folded + out, routed, _, _ = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = None): + """Attention of layer ``i`` up to its o_proj (no reduction). ``push`` is the attention boundary's exchange the + o_proj may write into. Returns (o_proj output, pushed).""" + core = self.core + if i not in core.sparse_layers: + return self._dense_attention(i, position_ids, hidden, step, push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _dense_attention(self, i: int, position_ids, hidden, step, push): + """Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v: an FP8 cache, or the FP8 pool the dense layers sit + in within a hybrid NVFP4 cache on a decode step), the attention core, o_proj. + + On a covered pure generation step (``_fullctx_plan``) the attention is this target's PrimTS split-KV decode + over the layer's full causal context, and the norm/RoPE op also stores the step's new K/V rows into the + layer's paged pool and the query compact, as the decode reads it (``DENSE_KV_INSERT_FUSED``; the bytes + ``write_msa_phase_kv`` stores). Otherwise the MSA backend.""" + core = self.core + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + plan = self._fullctx_plan(i, qkv.device, step) + kv_dest = None + if ( + plan is not None + and fp8 + and DENSE_KV_INSERT_FUSED + and _dense_fp8_op_applies(qkv, position_ids) + and hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + ): + kv_dest = ( + step["kv_cache_manager"].get_buffers(i, kv_layout="HND"), + step["msa_out_cache_loc"], + plan[1], + torch.empty( + (qkv.shape[0], core.heads, core.head_dim), + dtype=torch.float8_e4m3fn, + device=qkv.device, + ), + ) + if fp8: + fused = _dense_qk_norm_rope_fp8(core, i, qkv, position_ids, kv_dest) + if kv_dest is not None: + taken(DENSE_KV_INSERT_PATH, kv_dest[2]) + else: + fused = _qk_norm_rope(core, i, qkv, position_ids, False) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + if kv_dest is not None: + q = kv_dest[3].view(q.shape[0], -1) + o = None + if plan is not None: + o = self._fullctx_decode(i, q, k, v, step, plan, kv_written=kv_dest is not None) + if o is None: + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _fullctx_plan(self, i: int, device: torch.device, step): + """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs + on ``kernels/m3_fullctx_decode`` (``FULLCTX_DECODE``): a covered (query rows per request, requests) shape over + the layer's FP8 pool. None otherwise (the MSA backend runs). + + Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 (29.3), (1, 1) + 15.7 (28.0), (4, 2) 24.5 (29.3), (1, 2) 20.2 (28.1); the stock split is capped at 16 cluster CTAs per (request, + KV head), i.e. 32 of 212 CTAs at one request. With the BF16 split-KV merge (``FULLCTX_DECODE_OWN_REDUCER``): + (1, 1) 14.5 vs 18.5, (4, 1) 17.3 vs 21.3, (1, 2) 18.5 vs 23.1, (4, 2) 23.3 vs 27.5.""" + path = FULLCTX_DECODE.path + span = step["msa_decode_span"] + if span is None or span.row_first != 0: + uncovered(path, "no decode span starting at row 0") + return None + nreq = int(step["num_seqs"]) + qlen = int(span.query_len) + if not FULLCTX_DECODE.covers((qlen, nreq)): + uncovered(path, f"{qlen} query rows x {nreq} requests is not a covered shape") + return None + if int(step["max_seq_len"]) > FULLCTX_DECODE_MAX_KV_LEN: + uncovered(path, "the cache length bound is beyond the plans' K/V capacity") + return None + num_tokens = nreq * qlen + pool, _ = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + if pool.dtype != torch.float8_e4m3fn: + missed(path, num_tokens, "the dense layer's KV pool is not FP8") + return None + from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + + own_reducer = FULLCTX_DECODE_OWN_REDUCER + if own_reducer and not _fullctx_decode_kernel.has_own_reducer(): + missed(path, num_tokens, "the BF16 split-KV merge op is not built") + return None + # Compiled before warmup (``prebuild_kernels``); the forward never compiles one. + plan = _fullctx_decode_kernel.plan_for(device, qlen, nreq, own_reducer) + if plan is None: + missed(path, num_tokens, "the PrimTS decode plan is not prebuilt") + return None + return plan, num_tokens + + def _fullctx_decode(self, i: int, q, k, v, step, plan_tokens, *, kv_written: bool): + """Dense layer ``i``'s decode attention on the PrimTS plan of ``_fullctx_plan``: the step's new K/V go into the + paged cache as the MSA backend writes them (``write_msa_phase_kv``; already there when ``kv_written``), then + every query row attends its request's full causal context.""" + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + write_msa_phase_kv, + ) + + from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + + core = self.core + plan, num_tokens = plan_tokens + nreq = int(step["num_seqs"]) + if not kv_written: + write_msa_phase_kv( + core.attn[i], + k[:num_tokens], + v[:num_tokens], + step["metadata"], + AttentionInputType.mixed, + token_offset=0, + ) + pool, subpages = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + k_pages, v_pages = _fullctx_decode_kernel.kv_page_views(pool, subpages) + _record("m3_fullctx_decode", num_tokens, nreq, kv_written) + plan.run( + fp8_e4m3_saturate(q[:num_tokens].reshape(num_tokens, core.heads, core.head_dim)), + k_pages, + v_pages, + step["msa_block_table"][:nreq], + step["msa_seq_lens_cuda"][:nreq], + core.sm_scale(i), + output[:num_tokens], + ) + taken(FULLCTX_DECODE.path, num_tokens) + return output + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step, push): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection and + its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and the FP8 + indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the o_proj. A + layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): a sparse + layer whose caches the producer wrote runs this target's selection and sparse decode when the generation span + starts at row 0; everything else the MSA backend (``m3_msa_attention``: the KV write unless a producer wrote the + caches, the sparse layers' top-k block selection, the FMHA).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + span = step["msa_decode_span"] + sparse = i in core.sparse_layers + if sparse and k_live is None and span is not None and span.row_first == 0: + self._decode_sparse_core(i, q_live, idx_q_live, step, out_live) + return output + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + def _index_scores(self, i: int, idx_q, step, num_tokens: int, qlen: int): + """The per-index-head block maxima of a non-identity decode step: this target's CuTe DSL scorer at the request + counts ``INDEX_SCORE`` covers (each count's CTAs per request and early issue, ``index_score_schedule``), the + in-tree CuTe DSL scorer beyond them. Rubin TP2 (2 index heads), 165,521 tokens, behind a 58 MB PDL producer, + per layer: 1 request 3.27 us at 1536 CTAs (in-tree 256-CTA scorer 6.07); 2 requests 4.97 us at 768 (in-tree + 6.71).""" + config = self.core.attn[i].m3_config + rows = int(step["num_seqs"]) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = step["msa_idx_k_cache"](i) + max_score = step["msa_proxy_max_score_view"]( + config.num_index_heads, step["msa_worst_case_max_k_tiles"], num_tokens + ) + if INDEX_SCORE.covers(rows): + split_k, early = index_score_schedule(rows) + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(split_k), + bool(early), + ) + if key not in _index_score_kernel.IndexScoreRunner.kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + missed( + INDEX_SCORE_PATH, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score", key) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + split_k, + early, + ) + taken(INDEX_SCORE_PATH, rows) + else: + uncovered( + INDEX_SCORE_PATH, + f"{rows} requests, past {INDEX_SCORE_MAX_REQUESTS}: the in-tree CuTe DSL scorer", + ) + _in_tree_scorer() # registers the op + _record("cute_dsl_minimax_m3_index_decode_score", rows, qlen) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + ) + return max_score + + def _select(self, i: int, max_score, step, num_tokens: int): + """The top-16 table of a non-identity decode step from the block maxima reduced to the KV head: the long-row + select (the 384-thread kernel, ``SELECT_LONG_DSL``, when compiled) on rows of ``LONG_SELECT_MIN_BLOCKS.. + LONG_SELECT_MAX_BLOCKS`` blocks, else the built-in selector (``m3_select_blocks``).""" + config = self.core.attn[i].m3_config + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + group = config.num_index_heads // config.num_kv_heads + max_score_kv = ( + max_score.view(config.num_kv_heads, group, max_score.shape[1], max_score.shape[2]).amax( + dim=1 + ) + if group > 1 + else max_score + ) + from .kernels import m3_select_long_dsl as _select_long_dsl_kernel + + num_blocks = int(max_score_kv.shape[1]) + table = None + dsl = False + if long_select_covers(num_blocks): + early_trigger = num_tokens <= SELECT_EARLY_MAX_ROWS + if LONG_SELECT_DSL.keys: + # Compiled before warmup (prebuild_kernels); a missing kernel falls back to the 1024-thread one. + key = _select_long_dsl_kernel.kernel_key( + get_env_enable_pdl(), early_trigger, SELECT_LONG_TRIGGER_PHASE + ) + dsl = key in _select_long_dsl_kernel.kernel_cache + if not dsl: + missed( + LONG_SELECT_DSL.path, + num_tokens, + f"the 384-thread select {key} was not compiled before warmup", + ) + _record("long_select", num_blocks, num_tokens, dsl) + table = long_select_blocks( + max_score_kv, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + use_pdl=get_env_enable_pdl(), + early_trigger=early_trigger, + trigger_phase=SELECT_LONG_TRIGGER_PHASE, + dsl=dsl, + ) + if table is not None: + taken((LONG_SELECT_DSL if dsl else LONG_SELECT).path, num_tokens) + else: + long_select_fallback(num_blocks, num_tokens) + _record("m3_select_blocks", num_blocks, num_tokens) + table = m3_select_blocks( + max_score_kv, + MSA_REQUIRED_TOPK, + n_valid_blocks, + config.init_blocks, + config.local_blocks, + True, + ) + return table + + def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step of a sparse layer whose caches + the producer already wrote. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise the + index scorer writes the per-index-head block maxima, reduced to the KV head, and the long-row selector takes the + top 16 (the built-in selector beyond its shapes). On an NVFP4 layer the decode (``m3_sparse_decode``: the + long-context decode or the Triton producer + merge) prefetches into L2, before its grid wait, the K/V pages of + the blocks the same layer selected one step back (``SPARSE_DECODE_LC_PREFETCH_SELF_KV``: 65-68% of a layer's + selected blocks recur one step later; same layer, not the next one as at TP4: the MoE between two sparse layers + streams ~250 MB through the 126 MB L2, so an earlier issue would be evicted; hints only). + """ + from .kernels import m3_attn_lc as _attn_lc_kernel + + core = self.core + backend = core.attn[i] + num_tokens = int(q.shape[0]) + rows = int(step["num_seqs"]) + qlen = int(step["msa_decode_span"].query_len) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + num_kv_heads = int(backend.indexer.config.num_kv_heads) + # The index scorers are compiled and loaded by prebuild_kernels, so no step compiles one and the identity can + # be taken from the first decode step. + identity = step["msa_short_ok"] and step["msa_short_table"] is not None + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + else: + max_score = self._index_scores(i, idx_q, step, num_tokens, qlen) + topk = self._select(i, max_score, step, num_tokens).permute(1, 0, 2) + manager = step["kv_cache_manager"] + k_paged, v_paged = msa_paged_kv(manager, i) + nvfp4 = {} + if manager.is_nvfp4_layer(i): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(backend, core.kv_scales[i]) + scales = manager.get_block_scale_buffers(i, "HND") + nvfp4 = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + # The long-context decode (if it runs) takes the layer's N1 id buffer; any other route leaves the slot unused, + # and it is cleared after the call either way. + n1_ids = core.n1_ids.get(i) + if n1_ids is not None and nvfp4 and num_tokens <= N1_MAX_ROWS: + _attn_lc_kernel.N1_IDS.ids = n1_ids + _record("m3_sparse_decode", num_tokens, qlen, identity) + try: + m3_sparse_decode( + q.view(num_tokens, core.heads, core.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=core.sm_scale(i), + output=output.view(num_tokens, core.heads, core.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4, + ) + finally: + _attn_lc_kernel.N1_IDS.ids = None + + +_ENGAGED: set = set() + + +def _log_engaged(path: str, mapping) -> None: + """One info line per fast path and process (every GEN rank logs its own): the run's engagement evidence.""" + if path not in _ENGAGED: + _ENGAGED.add(path) + rank = getattr(mapping, "tp_rank", "?") if mapping is not None else "?" + logger.info(f"[M3 sm_107 tp2] fast path engaged: {path} (tp rank {rank})") + + +class M3SpecStepFold: + """The target's speculative step-input staging (the in-tree ``SpecStepFold`` protocol). + + Between two decode graphs of a one-model Eagle3 step, the spec sampler stores the step's outputs into its per-slot + store and the engine gathers them back into its flat input buffers for the next step. This object does both in + one launch each (``m3_spec_tail.slot_scatter`` / ``step_gather``): integer copies only, the same values the + engine's and the sampler's own copies write. + + The forward of a decode-only step also does both at its tail (``store_in_forward``, one ``m3_fold_tail`` launch + inside the step graph, this target's C++ op): it stores the step's outputs for the rows ``begin_step`` armed and + stages the next step's inputs for those rows in that order. The arming is a device-side row count and slot list + staged before the replay, so a captured graph whose step was not armed writes nothing. The sampler then skips its + store for exactly the armed rows (``sampler_rows_stored``), and the next step skips its gathers when its rows are + those rows (``begin_step``). The sampler copies the store to the host on its D2H side stream and records + ``store_copies_done`` after the copies; the forward's tail waits on it before it rewrites the store. + + * ``store`` -- the spec sampler's per-slot store (``bind_sampler_store``). + * ``step_inputs`` -- the engine's (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) device buffers + (``bind_engine_buffers``). + """ + + # Read by the shell's forward (the tail runs only when the fold stores in the forward). + in_forward = True + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[tuple] = None + # Device arming read by the forward's tail: ctrl[0] = armed rows (0 = write nothing), slots[:rows]. + self._ctrl: Optional[torch.Tensor] = None + self._slots: Optional[torch.Tensor] = None + # The rows (sequence slots) the next forward stores and stages for, as begin_step armed them. + self._armed: Optional[tuple] = None + # The rows the previous forward staged the next step's inputs for (confirmed by sampler_rows_stored). + self._staged_for: Optional[tuple] = None + # Output buffers the tail ran on: the captured graphs' static outputs (by address) and the last eager step's. + self._graph_outputs: set = set() + self._eager_outputs: Optional[torch.Tensor] = None + self._logged = False + # The (external) event the sampler records on its D2H side stream after copying the store to the host, waited + # on by the forward's tail before it rewrites the store (see SpecSamplerBase._copy_outputs_to_host). + self.store_copies_done: Optional[torch.cuda.Event] = None + + def bind_engine_buffers( + self, input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ): + self.step_inputs = (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) + # Recorded once now so a captured wait has a completed record to wait on before the first copy. + self.store_copies_done = torch.cuda.Event(external=True) + self.store_copies_done.record() + logger.info( + "[M3 host] spec store host copies on the D2H side stream, awaited by the forward's tail" + ) + logger.info("[M3 host] spec fold tail: one C++ launch (no Triton) at the forward's tail") + self._ctrl = torch.zeros(1, dtype=torch.int32, device=input_ids.device) + self._slots = torch.zeros(max(max_batch_size, 1), dtype=torch.long, device=input_ids.device) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + if self._ctrl is None: + return False + rows = tuple(seq_slots) if seq_slots is not None else None + staged = rows is not None and rows == self._staged_for + self._staged_for = None + armed = rows if rows is not None and 0 < len(rows) <= self._slots.numel() else None + # Read only by this step's forward, so both uploads may wait for the step-input pack's flush. + if armed is not None: + slots = torch.tensor(armed, dtype=torch.long, pin_memory=prefer_pinned()) + copy_h2d_if_changed(self._slots[: len(armed)], slots, pack=True) + ctrl = torch.tensor( + [len(armed) if armed is not None else 0], dtype=torch.int32, pin_memory=prefer_pinned() + ) + copy_h2d_if_changed(self._ctrl, ctrl, pack=True) + self._armed = armed + return staged + + # The launches the store wait may precede. store_copies_done is an external event: under graph capture its wait + # is a node of its own, so the kernel launched right after it loses its programmatic (PDL) edge and with it the + # overlap of its pre-wait part. Only a launch that waits for its predecessor's completion anyway may follow it: + # the fold tail (a plain launch, no griddepcontrol). Waiting before the last draft_step instead cost +13 / +17 + # us/step. + STORE_WAIT_BEFORE = frozenset({"fold_tail"}) + + def _wait_store_copies(self, before: str) -> None: + """Order the store's rewrite after the sampler's host copies of it; ``before`` names the launch that follows. + Under graph capture a launch outside STORE_WAIT_BEFORE fails the capture.""" + if self.store_copies_done is None: + return + if before not in self.STORE_WAIT_BEFORE and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"M3SpecStepFold: the store wait would precede {before!r}, a launch outside " + f"{sorted(self.STORE_WAIT_BEFORE)}; an external-event wait node drops the next kernel's PDL edge" + ) + torch.cuda.current_stream().wait_event(self.store_copies_done) + + def store_in_forward(self, outputs: dict) -> None: + """At the forward's tail: store the armed rows' outputs and stage the next step's inputs (one launch).""" + if self._ctrl is None or self.store is None or self.step_inputs is None: + return + store = self.store + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + new_tokens = outputs.get("new_tokens") + next_new_tokens = outputs.get("next_new_tokens") + new_tokens_lens = outputs.get("new_tokens_lens") + next_draft_tokens = outputs.get("next_draft_tokens") + if any( + t is None or not isinstance(t, torch.Tensor) + for t in (new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens) + ): + return + rows = new_tokens.shape[0] + # The previous step's host copies of the store run on the sampler's side stream. + self._wait_store_copies(before="fold_tail") + _record("m3_fold_tail", rows) + if not m3_fold_tail.fold_tail( + self._ctrl, + self._slots, + min(rows, self._slots.numel()), + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + next_draft_tokens, + store.next_draft_tokens, + new_tokens_lens, + store.new_tokens_lens, + input_ids, + draft_tokens, + position_id_offsets, + kv_lens_offsets, + next_new_tokens.shape[1], + cpp=True, + ): + return + if torch.cuda.is_current_stream_capturing(): + self._graph_outputs.add(new_tokens.data_ptr()) + else: + self._eager_outputs = new_tokens + + def gather( + self, + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + if self.step_inputs is None: + return False + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + return m3_spec_tail.step_gather( + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len, + tokens_per_gen_step, + draft_token_buffer_width, + input_ids, + num_tokens, + draft_tokens, + num_draft_tokens, + position_id_offsets, + kv_lens_offsets, + first_previous_row, + ) + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list) -> bool: + armed, self._armed = self._armed, None + new_tokens = outputs.get("new_tokens") + ran = isinstance(new_tokens, torch.Tensor) and ( + new_tokens is self._eager_outputs or new_tokens.data_ptr() in self._graph_outputs + ) + self._eager_outputs = None + stored = ( + ran + and armed is not None + and num_skip == 0 + and tuple(seq_slots) == armed + and new_tokens.shape[0] == len(armed) + ) + self._staged_for = armed if stored else None + if stored and not self._logged: + self._logged = True + logger.info("[M3 host] spec step stored and next step staged in the forward tail") + return stored + + def scatter( + self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) -> bool: + store = self.store + if store is None: + return False + return m3_spec_tail.slot_scatter( + seq_slots, + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + new_tokens_lens, + store.new_tokens_lens, + next_draft_tokens, + store.next_draft_tokens, + ) + + +class _DraftForward: + """``draft_model.model`` for the in-tree worker: the drafter layer's forward, called as + ``model(input_ids=, position_ids=, hidden_states=, attn_metadata=, spec_metadata=, all_rank_num_tokens=)`` and + returning (hidden, hidden to carry to the next pass). ``d2t`` None: the draft vocabulary is the target's.""" + + d2t = None + + def __init__(self, draft: "M3Eagle3Draft") -> None: + self._draft = draft + + def __call__(self, **kwargs): + return self._draft.forward_layer(**kwargs) + + +# The drafter's projections by weight-table name, for the low-M linears' engagement and prebuild. +_DRAFT_LINEARS = ("fc", "qkv", "o", "gu", "dn") + + +class M3Eagle3Draft(nn.Module): + """The one-layer BF16 Eagle3 drafter (``speculative_model``): its weights (``weights.DRAFT_WEIGHTS``), its + vocab-parallel BF16 head and logits processor, the drafter layer's TRTLLM attention backend (layer index = the + target's layer count, FP8 KV) and its TP all-reduce. + + What the worker reads: ``model`` (the layer forward), ``apply_eagle3_fc`` (the per-layer fc_norms + fc over the + captured target hidden states), ``lm_head`` and ``logits_processor``; the fused draft loop reads the layer's parts + (``draft_input_weights``, ``attention``, ``mlp``, ``ar_norm``). The drafter embeds with the target's table. + + The BF16 projections (fc, qkv, o_proj, gate_up, down; not the head) run FlashInfer's fixed-tactic low-M kernels + (``ops/gemm/m3_draft_lowm``, the TP8 port) at the rows ``DRAFT_LOWM`` covers, compiled before the warmup + (``prebuild_linears``); every other call is the in-tree cuBLAS GEMM. Rubin TP2 (mclk locked, 4 rotating sets) at + 1 row: qkv 13.18 -> 9.71, o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. The + split-K ('s') tactics run this target's copy of FlashInfer's kernel (``kernels/m3_lowm_splitk``, bit for bit the + same output); with ``DRAFT_AR_PUSH`` the o_proj and down projection of the fused draft loop also push their rows + into the drafter boundary all-reduce's exchange (``arm_push`` / ``take_pushed``), as do the direct ('d') 1-row + linears through this target's copy of the direct kernel (``kernels/m3_lowm_direct``; o_proj 9.22 -> 8.09 us, + down 13.81 -> 12.72 us per 1-row boundary). + """ + + def __init__( + self, + core: MiniMaxM3Tp2ModelingV2, + draft_config: ModelConfig, + start_layer_idx: int, + capture_layers: int, + ) -> None: + super().__init__() + cfg = draft_config.pretrained_config + _ensure_draft_vocab_size(cfg) + eagle = getattr(cfg, "eagle_config", None) or {} + mapping = draft_config.mapping + self.draft_config = draft_config + # The engine's draft weight mapper binds to these two. + self.model_config = draft_config + self.config = cfg + self.capture_layers = int(capture_layers) + self.layer_idx = int(start_layer_idx) + self.dtype = cfg.torch_dtype + self.hidden = int(cfg.hidden_size) + self.hidden_size_in = int(getattr(cfg, "target_hidden_size", cfg.hidden_size)) + self.rms_eps = float(cfg.rms_norm_eps) + self.return_hidden_post_norm = bool( + eagle.get("return_hidden_post_norm", False) or getattr(cfg, "norm_output", False) + ) + if ( + int(cfg.num_hidden_layers) != 1 + or eagle.get("eh_proj_before_attn", False) + or eagle.get("norm_before_fc", False) + or not getattr(cfg, "fc_norm", False) + or getattr(cfg, "kv_lora_rank", None) + or self.capture_layers <= 1 + or self.hidden_size_in != self.hidden + or cfg.draft_vocab_size != cfg.vocab_size + or getattr(cfg, "attention_bias", False) + or getattr(cfg, "mlp_bias", False) + or getattr(cfg, "bias", False) + or getattr(cfg, "tie_word_embeddings", False) + ): + raise ValueError( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle3 drafter with per-layer " + "fc_norms over " + "several captured layers, the target's hidden size and vocabulary and no biases" + ) + tp = mapping.tp_size + self.heads = cfg.num_attention_heads // tp + self.kv_heads = (cfg.num_key_value_heads + tp - 1) // tp + self.head_dim = int( + getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + ) + self.q_size = self.heads * self.head_dim + self.kv_size = self.kv_heads * self.head_dim + self.w = _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + self.lm_head = LMHead( + cfg.draft_vocab_size, + cfg.hidden_size, + dtype=cfg.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=draft_config.lm_head_gather_output, + reduce_output=False, + use_custom_cublas_mm=False, + quant_config=None, + ) + self.logits_processor = LogitsProcessor() + self.allreduce = AllReduce( + mapping=mapping, strategy=draft_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + quant_config = draft_config.get_quant_config() + self.attn = create_attention( + draft_config.attn_backend, + self.layer_idx, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, rope=RopeParams.from_config(cfg) + ), + quant_config=quant_config, + skip_create_weights_in_init=draft_config.skip_create_weights_in_init, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): + raise ValueError( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + self._bcg_key = f"m3_tp2_draft_{id(self)}" + _BCG_DRAFTS[self._bcg_key] = self + # The drafter embeds with the target's table, read off the core on every call (the engine materializes the + # core's weights after construction). + self._core = [core] + self.model = _DraftForward(self) + # The low-M linears: (N, K, M, PDL) of every call prebuild_linears compiled, and (N, K, M) of every call that + # logged its engagement line. + self._lowm_built: set = set() + self._lowm_engaged: set = set() + # DRAFT_AR_PUSH: the weight whose next low-M call also pushes its rows into a boundary all-reduce (data + # pointer, the bucket's multicast words, its flags, rank, world, max tokens), set by arm_push; and the data + # pointer of the output the last armed call pushed, 0 if none (read once by take_pushed). Host state set and + # read in one forward's issue order. + self._push_arm: Optional[tuple] = None + self._pushed = 0 + + @property + def embed_table(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def load_weights_from_target_model(self, target_model: nn.Module) -> None: + """Nothing to alias: the drafter reads the target's embedding table through the core and owns its head.""" + + # ------------------------------------------------------------------------------------------- the projections + + def linear(self, x: torch.Tensor, key: str) -> torch.Tensor: + """``x @ w[key]^T``: the low-M kernel where ``DRAFT_LOWM`` covers the call, else the in-tree cuBLAS GEMM.""" + w = self.w[key] + out = self._lowm_linear(w, x) + return out if out is not None else m3_bf16_linear(x, w) + + def _lowm_linear(self, weight: torch.Tensor, x: torch.Tensor) -> Optional[torch.Tensor]: + """x [M, K] @ weight^T through the low-M kernel when ``DRAFT_LOWM`` covers it; None for the in-tree GEMM. + A projection none of whose row counts is covered never reaches the low-M checks.""" + n, k = weight.shape + if x.dim() != 2 or not any(nk == (n, k) for nk, _ in DRAFT_LOWM.keys): + return None + m = int(x.shape[0]) + if not DRAFT_LOWM.covers(((n, k), m)): + uncovered(DRAFT_LOWM.path, f"the drafter's {n}x{k} linear past its tactic table's rows") + return None + if not x.is_contiguous(): + # The first pass's fc may read a column window of the captured hidden states; the in-tree GEMM takes + # strides. + uncovered(DRAFT_LOWM.path, "a strided activation (the in-tree GEMM reads it in place)") + return None + use_pdl = get_env_enable_pdl() + if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + missed(DRAFT_LOWM.path, m, "the activation or weight is not BF16") + return None + if (n, k, m, use_pdl) not in self._lowm_built: + missed( + DRAFT_LOWM.path, m, f"the {n}x{k} kernel at {m} rows was not compiled before warmup" + ) + return None + kind, tactic = tactic_for(n, k, m) + if (n, k, m) not in self._lowm_engaged: + self._lowm_engaged.add((n, k, m)) + logger.info( + f"[M3 sm_107 tp2] fast path engaged: m3_draft_lowm {n}x{k} at {m} rows (device {x.device.index})" + ) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + ar = None + if ( + self._push_arm is not None + and self._push_arm[0] == weight.data_ptr() + and ( + (kind == "s" and DRAFT_LOWM_OWN_SPLITK) + or (kind == "d" and DRAFT_AR_PUSH_DIRECT and tactic[1] % 2 == 0) + ) + ): + ar = self._push_arm[1:] + self._pushed = out.data_ptr() + _record("m3_draft_lowm", n, k, m, kind, ar is not None) + return m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, DRAFT_LOWM_OWN_SPLITK, ar) + + def prebuild_linears(self) -> int: + """Compile (and run once) every covered row count of every drafter projection, outside capture, on the + current stream, before the executor's warmup. Returns the number of calls compiled.""" + from .kernels.m3_lowm_direct import compile_direct_push + from .kernels.m3_lowm_splitk import compile_push + + use_pdl = get_env_enable_pdl() + built = 0 + shapes = {tuple(self.w[key].shape): self.w[key] for key in _DRAFT_LINEARS} + for (n, k), weight in sorted(shapes.items()): + if weight.dtype != torch.bfloat16: + continue + for _, m in sorted(key for key in DRAFT_LOWM.keys if key[0] == (n, k)): + if (n, k, m, use_pdl) in self._lowm_built: + continue + kind, tactic = tactic_for(n, k, m) + x = torch.zeros(m, k, dtype=torch.bfloat16, device=weight.device) + out = torch.empty(m, n, dtype=torch.bfloat16, device=weight.device) + m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, DRAFT_LOWM_OWN_SPLITK) + if DRAFT_AR_PUSH and DRAFT_LOWM_OWN_SPLITK and kind == "s": + # The push variant a boundary all-reduce may arm (o_proj / down): compiled only, never run here + # (its pushes need the all-reduce's exchange and a matching call). + compile_push(x, weight.t(), out, use_pdl, tactic) + if DRAFT_AR_PUSH and DRAFT_AR_PUSH_DIRECT and kind == "d" and tactic[1] % 2 == 0: + # The direct kernel's push variant (kernels/m3_lowm_direct), compiled only. + compile_direct_push(x, weight.t(), out, use_pdl, tactic) + self._lowm_built.add((n, k, m, use_pdl)) + built += 1 + return built + + def arm_push(self, key: str, workspace, rows: int) -> None: + """Arm ``w[key]``'s next covered low-M call to push into ``workspace``'s bucket for ``rows``.""" + bucket = m3_mc_allreduce._bucket(workspace, rows) + flags = None if bucket is None else bucket.flags_for(8) + self._pushed = 0 + self._push_arm = None + if bucket is not None and flags is not None and bucket.max_tokens == 8: + self._push_arm = ( + self.w[key].data_ptr(), + bucket.mc.data_ptr(), + flags.data_ptr(), + bucket.rank, + bucket.world_size, + bucket.max_tokens, + ) + + def take_pushed(self, x: torch.Tensor) -> bool: + """Whether the armed call pushed its rows; disarms. ``x`` is what the boundary all-reduce is about to reduce: + it must be the pushing call's own output (the prepushed all-reduce reads only the pushed words, never ``x``), + so an operation between that call and the all-reduce raises instead of silently reducing the wrong rows.""" + pushed, self._pushed, self._push_arm = self._pushed, 0, None + if pushed and x.data_ptr() != pushed: + raise RuntimeError( + "the drafter's pushed rows are not the boundary all-reduce's input (an operation ran in between)" + ) + return bool(pushed) + + # ------------------------------------------------------------------------------------------------- the layer + + def fc(self, hs: torch.Tensor) -> torch.Tensor: + return self.linear(hs, "fc") + + def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The captured target hidden states [T, layers x H] to the drafter's input width: each layer's fc_norm, + then the fc (the in-tree ``Eagle3ForCausalLM.apply_eagle3_fc``).""" + hidden_states = hidden_states.to(self.dtype) + if hidden_states.shape[-1] != self.hidden: + chunks = hidden_states.chunk(self.capture_layers, dim=-1) + hidden_states = torch.cat( + [ + flashinfer_rmsnorm(c, self.w[f"fc_norm{i}"], self.rms_eps) + for i, c in enumerate(chunks) + ], + dim=-1, + ) + hidden_states = self.fc(hidden_states) + return hidden_states + + def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: + """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is + False, where the caller fuses the reduction). The backend's generation phase runs this target's full-context + decode with the fused RoPE + FP8 K/V append when the shell installed it (``kernels/m3_drafter_attn``).""" + qkv = self.linear(x, "qkv") + _record("m3_draft_attention", int(md.num_tokens)) + if is_in_breakable_cuda_graph(): + # A context step's draft pass is captured with the target's: the attention runs in an eager break, on the + # step's metadata (the drafter's: the runtime's draft swap installs it on the target's object). + if get_model_extra_attrs()["attention_metadata"]() is not md: + raise RuntimeError( + "MiniMax-M3 drafter under a breakable prefill CUDA graph: its metadata is not the step's" + ) + o = qkv.new_empty((qkv.shape[0], self.q_size), dtype=torch.bfloat16) + _draft_attention_eager(qkv, o, self._bcg_key) + else: + o = m3_draft_attention(self.attn, qkv, md) + out = self.linear(o, "o") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool, cpp_silu: bool = False) -> torch.Tensor: + """gate_up, the plain SwiGLU, down (all-reduced unless ``reduce`` is False). ``cpp_silu`` (the fused draft + loop): the SwiGLU on this target's C++ op (``ops/spec/m3_draft_silu``, minimaxM3Sm107Tp2/siluAndMul.cu), bit + for bit the in-tree Triton silu_and_mul in one launch either way (no Triton on the GEN decode path).""" + gu = self.linear(h, "gu") + act = None + if cpp_silu: + act = m3_draft_silu(gu) + if act is None: + uncovered( + DRAFT_SILU_PATH, "the drafter SwiGLU input layout (the Triton op runs it)" + ) + else: + taken(DRAFT_SILU_PATH, int(gu.shape[0]), "drafter MLP SwiGLU, one C++ launch") + if act is None: + act = m3_silu_and_mul(gu, None, None, None) + out = self.linear(act, "dn") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def ar_norm(self, x: torch.Tensor, residual: torch.Tensor, norm_key: str): + """One draft-layer boundary in the in-tree structure: the projection's plain all-reduce (BF16 out), then the + fused residual add + RMSNorm (FP32 add, norm of the FP32 sum). Returns (normed, residual sum).""" + return flashinfer_fused_add_rmsnorm( + m3_allreduce(x, self.allreduce, None), residual, self.w[norm_key], self.rms_eps + ) + + def forward_layer( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + all_rank_num_tokens: Optional[List[int]] = None, + ): + """The drafter layer (the in-tree ``Eagle3DraftModel.forward`` for one layer): the embedded tokens and the + fc'd hidden states normed and concatenated, attention, the post-attention norm with the residual, the MLP and + the final norm. Returns (normed output, the hidden state the next pass carries).""" + previous = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + embeds = m3_embedding(input_ids, self.embed_table).to(self.dtype) + residual = hidden_states + h = flashinfer_rmsnorm(hidden_states, self.w["hidden_norm"], self.rms_eps) + e = flashinfer_rmsnorm(embeds, self.w["in_norm"], self.rms_eps) + h = torch.cat([e, h], dim=-1) + h = self.attention(h, attn_metadata, reduce=True) + h, residual = flashinfer_fused_add_rmsnorm( + h, residual, self.w["post_norm"], self.rms_eps + ) + h = self.mlp(h, reduce=True) + spec_metadata.maybe_capture_hidden_states(self.layer_idx, h, residual) + h, to_save = flashinfer_fused_add_rmsnorm( + h, residual, self.w["final_norm"], self.rms_eps + ) + return (h, h) if self.return_hidden_post_norm else (h, to_save) + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous + + +def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, capture_norm) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "tp_size", 1) > 1 + ) + + +class _M3ShardDraftLogits(nn.Module): + """W2b: the drafter's logits processor on an in-tree draft step: this rank's BF16 vocab shard of the drafter head + (no all-gather, no fp32 cast), which the worker's greedy top-1 reads (``greedy_sample_draft_with_tp_gather``). + Rows the caller has not selected (``return_context_logits`` False) take the original processor.""" + + def __init__(self, original: nn.Module): + super().__init__() + self._m3_original = [original] # not a submodule: the target shares it + + def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits: bool = False): + if not return_context_logits: + return self._m3_original[0]( + hidden_states, lm_head, attn_metadata, return_context_logits + ) + return lm_head.apply_linear(hidden_states, None) + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify and fused draft loop, at the in-tree numerics. + + * Verify: on decode-only batches whose verify rows the fused kernels cover (``tail_eligible``), the target's BF16 + lm_head computes this rank's vocabulary shard of the logits (the in-tree head's own local GEMM, as two + half-vocabulary GEMMs on the die that holds each half at ``UGPU_HEAD``'s rows) and one kernel replaces the logits + all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. + * Draft: the drafter's first pass may take the per-layer fc_norms in one launch (or folded into the hidden-state + capture) and run its fc beside the target head GEMM (forked right after the last captured target layer at most + ``FC_EARLY_MAX_ROWS`` verify rows; after a sampled verify's head at most ``FC_AFTER_HEAD_MAX_ROWS``). On a + decode-only step ``DRAFT_STEP`` covers, the draft loop is fused: each pass's draft input in one launch, the + drafter layer's attention and MLP with their all-reduces fused with the residual add and the next RMSNorm (the + multicast all-reduce at ``DRAFT_MC_AR``'s rows, fed by the projection's push), and one kernel for the BF16 head, + the cross-rank argmax and the pass's bookkeeping. + + Every other batch takes the in-tree paths. Methods whose docstring says "Mirrors" copy the named in-tree method + with the fused branches inserted. + """ + + def __init__(self, spec_config, mapping, use_separate_draft_kv_cache: bool = False) -> None: + super().__init__( + spec_config, mapping=mapping, use_separate_draft_kv_cache=use_separate_draft_kv_cache + ) + # Per-step state, reset at the top of every _forward_impl: the first draft pass's gather rows when this step + # takes the fused first drafter inputs, and the fc output a fork computed beside the target head. + self._m3_first_gather_ids: Optional[torch.Tensor] = None + self._m3_fc_hidden: Optional[torch.Tensor] = None + self._m3_target_lm_head: Optional[nn.Module] = None + # (fc_norm0 weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple] = None + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + # Armed per step by arm_capture_norm: the drafter whose first fc is forked at the last captured target layer, + # and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + # The 12-bit coded copy of the target head shard the fused verify streams, built before the warmup. + # The die-placed target head (uGPU), built before the warmup, or None. + self._m3_ugpu_head = None + self._m3_force_vec_warm = False + # The streamed copy of the drafter's BF16 head shard the fused draft loop reads, built before the warmup, the + # head weight version it was built from, and the replaced copies a captured graph may still read. + self._m3_draft_head = None + self._m3_draft_head_version: Optional[int] = None + self._m3_draft_head_graveyard: list = [] + # Per step: the next step's input tokens [nreq, D + 1], column 0 written by the fused verify (or the seed) and + # columns 1.. by the fused draft passes; set as _m3_next_new_tokens once the fused draft loop filled it. + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + # Per step: the first draft inputs are the fused ones (accepted tokens as they lie, fc output of every row). + self._m3_fused_first_inputs = False + + def capture_target(self, spec_metadata, layer_id: int, num_tokens: int) -> Optional[tuple]: + """Where this step's capture norm writes layer ``layer_id``'s capture, for a boundary all-reduce that writes it + itself (CAPTURE_NORM_IN_BOUNDARY): (the capture buffer's [num_tokens, H] column block, the layer's fc_norm + weight, eps), or None (no capture norm armed, not a captured layer, or more rows than the buffer). The + boundary writes every row it reduces; rows past num_capture_tokens are unread.""" + capture_norm = self._m3_capture_norm + if capture_norm is None or spec_metadata.hidden_states is None: + return None + if num_tokens > spec_metadata.hidden_states.shape[0]: + return None + hidden_size = spec_metadata.hidden_size + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + out = spec_metadata.hidden_states[ + :num_tokens, i * hidden_size : (i + 1) * hidden_size + ] + return out, capture_norm[0][i], capture_norm[1] + return None + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None, written: bool = False + ) -> None: + """The Eagle3 capture with this step's capture norm (``written``: the layer's boundary all-reduce already wrote + it, see ``capture_target``); on a step armed for the early fc fork (``arm_capture_norm``), the last captured + layer then forks the drafter's first fc onto the aux stream, where it overlaps the remaining target layers + and the head (joined by ``_m3_target_accept``).""" + if not written: + _eagle3_capture(spec_metadata, layer_id, hidden_states, residual, self._m3_capture_norm) + drafter = self._m3_fc_early_drafter + layers = spec_metadata.layers_to_capture + if drafter is not None and layers and layer_id == layers[-1]: + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = self._m3_fc_fork( + spec_metadata, spec_metadata.num_capture_tokens, drafter + ) + + def prebuild(self, lm_head: nn.Module, draft: Optional[M3Eagle3Draft]) -> None: + """From the loaded weights, before the executor's warmup, on every rank: the die-placed copy of the target head + (uGPU), the fused draft loop's head copy and the drafter's low-M + linears.""" + start = time.perf_counter() + if self.prebuild_ugpu_head(lm_head): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] uGPU target head copy placed per die before warmup: %.1f s" + % (time.perf_counter() - start) + ) + start = time.perf_counter() + if self.prebuild_draft_head(draft): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + start = time.perf_counter() + built = draft.prebuild_linears() if draft is not None else 0 + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] drafter low-M GEMMs built before warmup: %d calls, %.1f s" + % (built, time.perf_counter() - start) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Mirrors Eagle3OneModelWorker._forward_impl, with the fused verify branch.""" + runtime_draft_len = spec_metadata.runtime_draft_len + # skip the draft forward if the runtime draft length is 0 + if runtime_draft_len == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._m3_first_gather_ids = None + self._m3_next_new = None + self._m3_next_new_tokens = None + self._m3_fused_first_inputs = False + if logits is None: + # tail_eligible() skipped the gathered target logits: verify from this rank's logits shard. + accepted_tokens, num_accepted_tokens, self._m3_first_gather_ids = ( + self._m3_target_accept( + hidden_states, + spec_metadata, + batch_size, + runtime_draft_len, + draft_model, + kv_lens=getattr(attn_metadata, "kv_lens_cuda", None), + ) + ) + else: + self._execute_guided_decoder_if_present(logits) + + # Sample and accept tokens. ``input_ids`` is required by the relaxed-acceptance path (scans for + # thinking-phase tokens); ignored otherwise. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, attn_metadata, spec_metadata + ) + if self.draft_only_eligible(attn_metadata, spec_metadata, draft_model): + self._m3_seed_first_pass(num_accepted_tokens, spec_metadata, batch_size) + + # Mamba hybrid models need state updates after token acceptance (a no-op on this target's cache manager). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None) + if self.sa_enhancer is not None and sa_manager is not None: + self.sa_enhancer.extend_and_prepare( + sa_manager=sa_manager, + request_ids=spec_metadata.request_ids, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_gens=num_gens, + num_contexts=num_contexts, + max_draft_len=runtime_draft_len, + ) + + if self._uses_external_shared_target_kv: + next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + ) + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here so the post-loop restore (below) + # can put attn_metadata back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + # prepare next new tokens to support overlap scheduler + if self._m3_next_new_tokens is not None: + # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). + next_new_tokens = self._m3_next_new_tokens + else: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + # Commit model-owned target state only after draft preparation has succeeded; SpecWorkerBase aborts the + # still-pending snapshots on any exception above this point. + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices[:batch_size], + num_contexts, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """The fused first draft inputs on a step that verified with the fused kernels (or was seeded for them), + else the in-tree ones.""" + fused = self._m3_first_drafter_inputs( + position_ids, + accepted_tokens, + input_ids.shape[0], + attn_metadata.num_contexts, + spec_metadata, + draft_model, + ) + self._m3_fused_first_inputs = fused is not None + if fused is not None: + fused.update(attn_metadata=attn_metadata, spec_metadata=spec_metadata) + return fused + return super().prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + def _m3_structural_ok(self, attn_metadata, spec_metadata) -> bool: + """The batch conditions shared by the fused verify and the fused first drafter inputs: a decode-only plain + Eagle3 step without attention DP, guided decoding, relaxed acceptance or a suffix-automaton enhancer, of at + most WIDE_VERIFY_MAX_TOKENS verify rows. Host state only, so the TP ranks agree.""" + if self.is_mtp_eagle or self.use_dynamic_tree: + return False + if self._uses_external_shared_target_kv or self.sa_enhancer is not None: + return False + if ( + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp + ): + return False + if getattr(self.spec_config, "use_relaxed_acceptance_for_thinking", False): + return False + if attn_metadata.num_contexts != 0 or attn_metadata.num_seqs == 0: + return False + draft_len = spec_metadata.runtime_draft_len + if draft_len == 0 or attn_metadata.num_seqs * (draft_len + 1) > WIDE_VERIFY_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Whether this step verifies with the fused kernels instead of the gathered target logits. + + Decode-only Eagle3 batches (``_m3_structural_ok``) over an unpadded vocab-sharded bias-free target head, all + greedy at a request count ``SHARD_ACCEPT`` covers, or sampled under the temperature + top-p contract (see + _m3_sampled_verify_ok) at a request count ``SAMPLED_ACCEPT`` covers, with nothing else reading the target + logits. Every input is host state shared by all TP ranks, so the ranks agree. + """ + if not self._m3_structural_ok(attn_metadata, spec_metadata): + return False + if getattr(spec_metadata, "enable_penalty", False): + return False + if not _sharded_head(lm_head) or getattr(lm_head, "bias", None) is not None: + return False + draft_len = spec_metadata.runtime_draft_len + if spec_metadata.is_all_greedy_sample: + key = (draft_len, attn_metadata.num_seqs) + if not SHARD_ACCEPT.covers(key): + uncovered( + SHARD_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + else: + if not self._m3_sampled_verify_ok(spec_metadata): + return False + key = (draft_len, attn_metadata.num_seqs) + if not SAMPLED_ACCEPT.covers(key): + uncovered( + SAMPLED_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + self._m3_target_lm_head = lm_head + _log_engaged( + SHARD_ACCEPT.path if spec_metadata.is_all_greedy_sample else SAMPLED_ACCEPT.path, + self.mapping, + ) + return True + + def draft_only_eligible(self, attn_metadata, spec_metadata, draft_model) -> bool: + """Whether a step verified by the in-tree path can still take the fused first drafter inputs: up to + FIRST_DRAFTER_INPUTS_MAX_REQUESTS requests, the structural conditions of tail_eligible, without the + target-sampling ones.""" + if attn_metadata.num_seqs > FIRST_DRAFTER_INPUTS_MAX_REQUESTS: + uncovered( + "m3_first_drafter_inputs", f"in-tree verify at {attn_metadata.num_seqs} requests" + ) + return False + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None + + def _m3_seed_first_pass(self, num_accepted_tokens, spec_metadata, batch_size: int) -> None: + """The first draft pass's gather rows (the last accepted row of each request) from the in-tree acceptance, as + _m3_target_accept leaves them.""" + k1 = spec_metadata.runtime_draft_len + 1 + n = num_accepted_tokens[:batch_size].to(torch.int64) + rows = spec_metadata.batch_indices_cuda[:batch_size].to(torch.int64) + self._m3_first_gather_ids = rows * k1 + n - 1 + + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], + ): + """The in-tree first draft inputs of a decode-only step, with the per-layer fc_norm RMSNorms of + ``apply_eagle3_fc`` in one kernel (unless the capture already applied them) before the fc; the input ids are + the accepted tokens as they lie. None to use the in-tree path.""" + if ( + self._m3_first_gather_ids is None + or self.is_mtp_eagle + or num_contexts != 0 + or draft_model is None + ): + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + ): + return None + if not self._m3_first_inputs_covered(num_tokens, spec_metadata): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not draft_model.w["fc_norm0"]: + if torch.cuda.is_current_stream_capturing(): + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), + ) + self._m3_fc_norm_w = cache + hidden = self._m3_fc_hidden + self._m3_fc_hidden = None + if hidden is None or hidden.shape[0] != num_tokens: + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + _record("m3_chunk_norm", num_tokens) + hs = m3_spec_tail.chunk_norm(hs, cache[1], draft_model.rms_eps) + _log_engaged(CHUNK_NORM.path, self.mapping) + hidden = draft_model.fc(hs) + input_ids = accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].reshape(-1) + _log_engaged(FUSED_FIRST_INPUTS.path, self.mapping) + return {"input_ids": input_ids, "position_ids": position_ids, "hidden_states": hidden} + + def _m3_first_inputs_covered(self, num_tokens: int, spec_metadata) -> bool: + """Whether the fused first drafter inputs of ``num_tokens`` verify rows are covered: their request count + (``FUSED_FIRST_INPUTS``), and the chunk norm's rows (``CHUNK_NORM``) unless the captures already applied the + fc_norm.""" + k1 = spec_metadata.runtime_draft_len + 1 + if num_tokens % k1 or not FUSED_FIRST_INPUTS.covers(num_tokens // k1): + uncovered( + FUSED_FIRST_INPUTS.path, + f"{num_tokens} verify rows at draft length {k1 - 1} (the in-tree first drafter inputs run)", + ) + return False + if self._m3_capture_norm is None and not CHUNK_NORM.covers(num_tokens): + uncovered( + CHUNK_NORM.path, f"{num_tokens} verify rows (the in-tree first drafter inputs run)" + ) + return False + return True + + def _m3_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + after: Optional[torch.cuda.Event] = None, + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (after the event + ``after`` when given, else after everything queued on the current stream); the result is picked up by + _m3_first_drafter_inputs. Returns the stream to join, or None when the fc path does not apply (the first + drafter inputs then compute it).""" + self._m3_fc_hidden = None + if draft_model is None: + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + or self._m3_fc_norm_w is None + or not self._m3_first_inputs_covered(num_tokens, spec_metadata) + ): + return None + aux = self._m3_aux_stream + if aux is None: + aux = self._m3_aux_stream = torch.cuda.Stream() + if after is not None: + aux.wait_event(after) + else: + aux.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(aux): + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + _record("m3_chunk_norm", num_tokens) + hs = m3_spec_tail.chunk_norm( + hs, self._m3_fc_norm_w[1], draft_model.rms_eps, use_pdl=False + ) + _log_engaged(CHUNK_NORM.path, self.mapping) + self._m3_fc_hidden = draft_model.fc(hs) + return aux + + def _m3_fc_join(self, fork: torch.cuda.Stream) -> None: + """Join the fc fork onto the current stream and keep its output's memory alive for that stream. + + _m3_fc_hidden is allocated in fork (aux) stream order and read on the current stream. Without the record, + dropping it returns its block to the aux stream's pool, where later aux work can take it while the current + stream's reads are still queued (eager warm-up / mixed steps). Nothing is recorded under graph capture, as in + maybe_execute_in_parallel: a captured graph relies on the next fork's aux.wait_stream(main) to order the + reuse.""" + main = torch.cuda.current_stream() + main.wait_stream(fork) + hidden = self._m3_fc_hidden + if ( + hidden is not None + and hidden.is_cuda + and not (torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()) + ): + hidden.record_stream(main) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model: Optional[M3Eagle3Draft] + ) -> None: + """Before every target forward: whether this step's Eagle3 captures apply the drafter's fc_norm themselves. + + On a fused-verify step of ``CAPTURE_NORM``'s rows whose first drafter inputs are the fused ones and whose + fc_norm weights are already stacked (by an earlier fused first pass), the capture writes + rmsnorm(h + r) * fc_norm[i] (one kernel per captured layer instead of add + copy, or the layer's boundary + all-reduce itself, ``capture_target``) and the fused first drafter inputs read it without their own + chunk_norm. A fused-verify step of at most ``FC_EARLY_MAX_ROWS`` verify rows is also armed to fork the + drafter's first fc at the last captured layer (``capture_hidden_states``). + """ + self._m3_capture_norm = None + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = None + if draft_model is None: + return + if not self.tail_eligible(attn_metadata, spec_metadata, lm_head): + return + rows = attn_metadata.num_seqs * (spec_metadata.runtime_draft_len + 1) + # The fork reads every captured column at the last captured layer, so a final-state capture, which is written + # after the model returns, keeps the fc at its stock place. + if ( + rows <= FC_EARLY_MAX_ROWS + and FINAL_HIDDEN_LAYER_ID not in spec_metadata.layers_to_capture + ): + self._m3_fc_early_drafter = draft_model + if not CAPTURE_NORM.covers(rows): + uncovered(CAPTURE_NORM.path, f"{rows} verify rows (the in-tree capture runs)") + return + # The normed captures are read only by the fused first drafter inputs; the in-tree ones norm again. + if not FUSED_FIRST_INPUTS.covers(attn_metadata.num_seqs): + return + fc_norm_w = self._m3_fc_norm_w + if ( + fc_norm_w is not None + and fc_norm_w[0] is draft_model.w["fc_norm0"] + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.rms_eps) + _log_engaged(CAPTURE_NORM.path, self.mapping) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata) -> bool: + """Whether sampled verify rows can use the fused sampled op: only under advanced_sampling_mode=no_topk, the + deploy-time guarantee that no row uses top-k or min-p (that op applies neither).""" + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + return mode is not None and mode.skips_top_k and not mode.skips_top_p + + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + kv_lens: Optional[torch.Tensor] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path. ``kv_lens``: the + requests' KV lengths at this verify (the per-request forced-acceptance key, FORCE_AL_PER_REQUEST).""" + force_total = 0 + force_vec = None + if self.force_num_accepted_tokens != 0.0: + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # The draw pool is built by an H2D copy and the draw kernels JIT-compile per block width, so both + # happen here, before the first captured step (a scratch counter keeps the real draw sequence). + self._ensure_force_accept_rng_state(hidden_states.device) + if not self._m3_force_vec_warm: + scratch = torch.zeros(1, dtype=torch.int64, device=hidden_states.device) + dummy = torch.empty(64, dtype=torch.int32, device=hidden_states.device) + for block in (16, 32, 64): + m3_force_vec( + self._force_accept_rng_pool, + scratch, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + m3_force_vec_req( + self._force_accept_rng_pool, + dummy, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + self._m3_force_vec_warm = True + if not is_warmup: + force_total = min(int(self.force_num_accepted_tokens) + 1, draft_len + 1) + force_vec = self._m3_force_vec( + batch_size, force_total, draft_len + 1, hidden_states.device, kv_lens + ) + draft_tokens = spec_metadata.draft_tokens.view(batch_size, draft_len) + if draft_tokens.dtype != torch.int32: + draft_tokens = draft_tokens.to(torch.int32) + # Column 0 receives the last accepted token, the next step's first input (the fused draft loop adds the rest). + self._m3_next_new = torch.empty( + (batch_size, draft_len + 1), dtype=torch.int32, device=hidden_states.device + ) + rows = hidden_states[: batch_size * (draft_len + 1)] + lm_head = self._m3_target_lm_head + batch_idx = spec_metadata.batch_indices_cuda[:batch_size] + # The Eagle3 fc on the aux stream: the fork made at the last captured layer, else one beside the head. A + # sampled verify's head (the in-tree GEMM up to FC_AFTER_HEAD_MAX_ROWS verify rows) holds + # every SM while it streams: its fc is forked after it instead (overlapping the sampled acceptance), where its + # many split-K CTAs do not delay the head's start. Rubin TP2 c1: the head started 17.4 us after its input was + # ready with the fc beside it. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + head_first = ( + fork is None + and rows.shape[0] <= FC_AFTER_HEAD_MAX_ROWS + and not spec_metadata.is_all_greedy_sample + ) + if fork is None and not head_first: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the head's own local GEMM (or its per-die split), whose + # all-gather the fused kernels replace (they exchange per-row records across ranks, not logits). + ugpu = self._m3_ugpu_head + if ugpu is not None and UGPU_HEAD.covers(rows.shape[0]): + if m3_ugpu_head_op.current(ugpu, lm_head.weight): + taken(UGPU_HEAD.path, rows.shape[0]) + _record("m3_ugpu_head", rows.shape[0]) + logits = m3_ugpu_head_op.m3_ugpu_head(rows, ugpu) + else: + missed( + UGPU_HEAD.path, + rows.shape[0], + "the placed head copy is not of the head's current weight (prebuild)", + ) + logits = lm_head.apply_linear(rows, None) + else: + logits = lm_head.apply_linear(rows, None) + head_done = None + if head_first: + head_done = torch.cuda.Event() + head_done.record() + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept_kernel.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = SHARD_ACCEPT.path + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept_kernel.sampled_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + spec_metadata.temperatures[:n_rows], + spec_metadata.top_ks[:n_rows], + spec_metadata.top_ps[:n_rows], + None, + seeds, + offsets, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = "spec_tail" + if head_done is not None: + # Enqueued after the sampled acceptance (which starts first) but waits on the head only. + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model, after=head_done) + if fork is not None: + self._m3_fc_join(fork) + if result is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, None, "the fused verify's exchange workspace") + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + return result + + def prebuild_ugpu_head(self, lm_head: nn.Module) -> bool: + """Build the die-placed copy of the target head shard for the uGPU verify head (``UGPU_HEAD``) before the + executor's warmup, on every rank (the module's weight is not changed): a verify of those rows computes the + logits shard as two half-vocabulary in-tree GEMMs, each on the green context of the die whose memory holds + that half, bit for bit the in-tree head's logits. Rubin TP2 head shard: 93.3 -> 79.5 us at 4 rows, 93.5 -> + 80.6 / 80.7 at 12 / 16; in-model the decode step is ~45 us shorter at one request and ~90 us at two. False + (nothing changed) when the head or device does not qualify (see m3_ugpu_head.prepare_head).""" + self._m3_ugpu_head = None + if not UGPU_HEAD.keys or not _sharded_head(lm_head): + return False + self._m3_ugpu_head = m3_ugpu_head_op.prepare_head(lm_head) + return self._m3_ugpu_head is not None + + def _m3_force_vec( + self, + batch_size: int, + base_total: int, + max_total: int, + device, + kv_lens: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + With ``kv_lens`` (FORCE_AL_PER_REQUEST): each request draws from the fixed-seed pool at its own key, its KV + length at this verify (its prompt plus what it has accepted so far), so its draws do not depend on the + server's history or on the other requests; every TP rank holds the same lengths, so the ranks agree. Without + lengths: exactly as _apply_force_accepted_tokens draws for a decode-only batch (one shared counter bump per + step, the (counter, slot) hash), the generic path's sequence. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: + return None + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + block = max(16, 1 << (batch_size - 1).bit_length()) + if FORCE_AL_PER_REQUEST and kv_lens is not None: + keys = kv_lens[:batch_size] + if keys.dtype != torch.int32: + keys = keys.to(torch.int32) + return m3_force_vec_req( + self._force_accept_rng_pool, + keys, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + # One launch: bump the shared counter and draw, with _apply_force_accepted_tokens' exact arithmetic. + return m3_force_vec( + self._force_accept_rng_pool, + self._force_accept_rng_counter, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[M3Eagle3Draft]) -> bool: + """Build the copy of the drafter's BF16 head shard that the fused draft loop streams (rebuilt only when the + head weight tensor or its version changed; a replaced copy is kept alive), on every rank. False when the + drafter is not one the loop serves.""" + old = self._m3_draft_head + if draft_model is None or not self._m3_drafter_supported(draft_model): + self._m3_draft_head = None + return False + weight = draft_model.lm_head.weight + # Idempotent: the same weight tensor at the same version keeps the built copy, so a repeated + # cache_derived_state never re-places it under a captured CUDA graph. + if ( + old is not None + and old.weight is weight + and self._m3_draft_head_version == weight._version + ): + return True + if old is not None: + # A rebuilt copy (new weight values): keep the old one alive, since a graph captured earlier may still + # read its (placed) memory until the graphs are recaptured. + self._m3_draft_head_graveyard.append(old) + self._m3_draft_head = m3_draft_step_kernel.prepare_head( + draft_model.lm_head, self.mapping.tp_rank + ) + self._m3_draft_head_version = weight._version + return True + + def _m3_drafter_supported(self, draft_model: M3Eagle3Draft) -> bool: + """The drafter the fused draft loop runs (``M3Eagle3Draft`` validates its shape at construction) over this + worker's tensor-parallel group with a bias-free vocab-sharded BF16 head without padding.""" + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if self._d2t is not None: + return False + lm_head = draft_model.lm_head + return _sharded_head(lm_head) and m3_draft_step_kernel.head_supported(lm_head) + + def _m3_draft_ar_norm( + self, + draft_model: M3Eagle3Draft, + x: torch.Tensor, + residual: torch.Tensor, + norm_key: str, + pushed: bool = False, + ): + """One drafter boundary of the fused draft loop: the multicast all-reduce + residual + RMSNorm with the + drafter's stock rounding (``DRAFT_MC_AR``: opts bit 16, normalizing the unrounded FP32 sum as FlashInfer + fused_add_rmsnorm does; Rubin TP2 10.0..11.0 -> 4.6..5.3 us per boundary), else the in-tree structure + (``M3Eagle3Draft.ar_norm``). ``pushed``: the boundary's projection already pushed its rows into the exchange + (prepushed 2). The choice depends on shapes and a collectively created workspace only, so the TP ranks + agree. Returns (normed, residual sum).""" + rows = x.shape[0] + path = DRAFT_MC_AR.path + if DRAFT_MC_AR.covers(rows): + workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) + out = None + if workspace is not None: + _record("m3_draft_mc_ar", rows, pushed) + out = m3_mc_allreduce.m3_mc_allreduce_norm( + x.contiguous(), + residual.contiguous(), + draft_model.w[norm_key], + draft_model.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + drafter_norm=True, + prepushed=2 if pushed else 0, + ) + if out is not None: + _log_engaged(path, self.mapping) + if pushed: + taken(DRAFT_AR_PUSH_PATH, rows, "the projection pushed its rows (prepushed 2)") + elif DRAFT_AR_PUSH: + uncovered( + DRAFT_AR_PUSH_PATH, + "the boundary's projection did not push (the AR pushes x)", + ) + return out + if pushed: + raise RuntimeError( + "the drafter's rows were pushed into a boundary all-reduce that did not run" + ) + missed(path, rows, "the multicast all-reduce workspace or operand layout") + return draft_model.ar_norm(x, residual, norm_key) + + def _m3_arm_draft_push(self, draft_model: M3Eagle3Draft, key: str, rows: int) -> None: + """DRAFT_AR_PUSH: let projection ``key``'s next low-M call push its rows into the boundary all-reduce that + ``_m3_draft_ar_norm`` then runs (the multicast path only; shapes and a collectively created workspace decide + it, so the TP ranks agree).""" + if not (DRAFT_AR_PUSH and DRAFT_MC_AR.covers(rows)): + return + workspace = m3_mc_allreduce.get_workspace(self.mapping, draft_model.w[key].shape[0]) + if workspace is not None: + draft_model.arm_push(key, workspace, rows) + + def _m3_fused_loop_ok( + self, attn_metadata, spec_metadata, batch_size: int, draft_model: Optional[M3Eagle3Draft] + ) -> bool: + """Whether this step drafts with the fused loop: a decode-only step the fused verify or the seed prepared + (``_m3_first_gather_ids``) with the fused first draft inputs, a supported drafter, a (draft length, requests) + ``DRAFT_STEP`` covers, the prebuilt head and the exchange workspace. Host state only, so the TP ranks + agree.""" + if ( + self._m3_first_gather_ids is None + or not self._m3_fused_first_inputs + or draft_model is None + ): + return False + path = DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + uncovered( + path, "the drafter or its attention metadata is not one the fused loop serves" + ) + return False + draft_len = spec_metadata.runtime_draft_len + if not DRAFT_STEP.covers((draft_len, batch_size)): + uncovered(path, f"draft length {draft_len} at {batch_size} requests") + return False + head = self._m3_draft_head + if head is None or head.weight is not draft_model.lm_head.weight: + missed( + path, + batch_size, + "the drafter head's streamed copy is not built for its current weight (prebuild)", + ) + return False + if m3_draft_step_kernel.workspace(self.mapping, head.wf.device) is None: + capture_missed(path, None, "the draft-step exchange workspace") + return False + _log_engaged(path, self.mapping) + return True + + def _forward_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ): + """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop (W2b: drafting from + the head's local shard where ``_m3_wide_draft_logits_ok``).""" + if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): + return self._m3_fused_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + # W2b: on a step the fused loop declines, the drafter head returns this rank's BF16 vocab shard (not the + # gathered fp32 logits), so the greedy top-1 runs the fused shard argmax (greedy_sample_draft_with_tp_gather). + swap = self._m3_wide_draft_logits_ok(spec_metadata, draft_model) + original = draft_model.logits_processor if swap else None + if swap: + draft_model.logits_processor = _M3ShardDraftLogits(original) + try: + return super()._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + finally: + if swap: + draft_model.logits_processor = original + + def _m3_wide_draft_logits_ok(self, spec_metadata, draft_model: Optional[nn.Module]) -> bool: + """Whether the in-tree draft loop may draft from the head's local shard (W2b): ``WIDE_SPEC_TAIL``, an Eagle3 + drafter with an unpadded vocab-sharded head under plain TP, argmax drafting (no rejection sampling) and no + guided decoder (which reads full logits). Host state only, so the TP ranks agree.""" + if not WIDE_SPEC_TAIL or draft_model is None or self.is_mtp_eagle: + return False + if self.guided_decoder is not None or spec_metadata.wants_advanced_draft_sampling: + return False + if self.mapping is None or self.mapping.enable_attention_dp: + return False + lm_head = getattr(draft_model, "lm_head", None) + return ( + lm_head is not None + and getattr(draft_model, "logits_processor", None) is not None + and _sharded_head(lm_head) + and getattr(lm_head, "bias", None) is None + ) + + def greedy_sample_draft_with_tp_gather( + self, logits: torch.Tensor, spec_metadata=None, mapping_lm_head_tp=None + ) -> torch.Tensor: + """W2b (``WIDE_SPEC_TAIL``): the in-tree draft loop's greedy top-1 over this rank's vocab shard of the drafter + logits, from the fused shard argmax + cross-rank key exchange (kernels/m3_shard_accept at draft length 0, at + most 48 rows per launch) instead of the per-rank max, the all-gather of (index, value) pairs and the pick. The + same tokens (keys order BF16 values exactly; ties take the lowest global index, as the pick over ranks and the + first-index max within a shard). Plain TP only (no LM-head TP in ADP); else, or when the logits are not a + unit-stride BF16 shard, or before the exchange workspace exists during a capture, the in-tree sampler.""" + tokens = None + if ( + WIDE_SPEC_TAIL + and mapping_lm_head_tp is None + and logits.dim() == 2 + and logits.dtype == torch.bfloat16 + and logits.stride(-1) == 1 + and self._draft_logits_are_sharded(logits, spec_metadata) + ): + m = logits.shape[0] + empty = torch.empty((m, 0), dtype=torch.int32, device=logits.device) + out = m3_shard_accept_kernel.shard_accept( + logits, + self.mapping.tp_rank * logits.shape[-1], + self.mapping, + empty, + m3_shard_accept_kernel._local_rows(m, logits.device), + 0, + ) + if out is not None: + tokens = out[0][:, 0] + _log_engaged("m3_draft_shard_top1", self.mapping) + if tokens is None: + return super().greedy_sample_draft_with_tp_gather( + logits, spec_metadata, mapping_lm_head_tp + ) + return tokens + + def _m3_fused_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model: M3Eagle3Draft, + draft_kv_cache_manager, + batch_size: int, + num_accepted_tokens, + ): + """Mirrors Eagle3OneModelWorker._forward_linear_draft_loop for a decode-only Eagle3 step, with each pass's + glue fused (see ``_m3_fused_loop_ok``): + + * the draft input -- the embedding lookup, the layer's input norm and hidden norm and their concat -- in one + launch (``m3_spec_tail.draft_input``; the second pass gathers the last accepted rows there); + * the drafter layer's attention and MLP (its SwiGLU on this target's C++ op), each followed by the all-reduce + with the residual add and the next RMSNorm (the post-attention norm, then the drafter's final norm): the + multicast boundary fed by the o_proj / down projection's push where covered (``_m3_draft_ar_norm``); + * one kernel for the BF16 head over the head rows (pass 0: the last accepted row of each request), the + cross-rank argmax and the bookkeeping the in-tree loop does between passes: the draft token, the next pass's + input id and position, the draft KV lengths (and pass 0's sequence lengths) (``m3_draft_step``). + + The host metadata transitions are the in-tree loop's. The drafts land in columns 1.. of the step's next-input + tokens (column 0 is the last accepted token), which ``_forward_impl`` returns as they lie. + """ + draft_len = spec_metadata.runtime_draft_len + table = draft_model.embed_table + dev = table.device + head = self._m3_draft_head + gather_ids = self._m3_first_gather_ids + ids = inputs["input_ids"] + hidden = inputs["hidden_states"] + pos = inputs["position_ids"] + if ids.dtype != torch.int32: + ids = ids.to(torch.int32) + if pos.dtype != torch.int32: + pos = pos.to(torch.int32) + num_acc = num_accepted_tokens[:batch_size] + if num_acc.dtype != torch.int32: + num_acc = num_acc.to(torch.int32) + next_new = self._m3_next_new + if next_new is None: + # Seeded from the in-tree acceptance: the last accepted token is the first-pass input at the gather row. + next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=dev) + next_new[:, 0] = ids[gather_ids] + draft_tokens = next_new[:, 1:] + w = draft_model.w + eps = draft_model.rms_eps + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) as md: + for i in range(draft_len): + x, residual = m3_spec_tail.draft_input( + ids, + table, + hidden, + gather_ids if i == 1 else None, + w["in_norm"], + eps, + w["hidden_norm"], + eps, + ) + rows = x.shape[0] + self._m3_arm_draft_push(draft_model, "o", rows) + h = draft_model.attention(x, md, reduce=False) + h, residual = self._m3_draft_ar_norm( + draft_model, h, residual, "post_norm", pushed=draft_model.take_pushed(h) + ) + self._m3_arm_draft_push(draft_model, "dn", rows) + h = draft_model.mlp(h, reduce=False, cpp_silu=True) + h, to_save = self._m3_draft_ar_norm( + draft_model, h, residual, "final_norm", pushed=draft_model.take_pushed(h) + ) + if draft_model.return_hidden_post_norm: + to_save = h + ids_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + pos_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + if not m3_draft_step_kernel.draft_step( + head, + h, + batch_size, + gather_ids if i == 0 else None, + self.mapping, + draft_tokens, + i, + ids_next, + pos, + pos_next, + md.kv_lens_cuda, + num_acc if i == 0 else None, + md._seq_lens_cuda if i == 0 else None, + draft_len, + ): + raise RuntimeError( + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" + ) + ids, pos, hidden = ids_next, pos_next, to_save + if i == 0: + # The draft step set the device sequence lengths to 1 and rolled the KV lengths back. + md._seq_lens[:batch_size].fill_(1) + md.on_update() + has_kv_cache = md.kv_cache_manager is not None + if has_kv_cache and hasattr(md, "host_request_types"): + md.host_request_types[: md.num_contexts].fill_(1) + if has_kv_cache: + md.num_contexts = 0 + self._prepare_flash_mla_generation_layout(md, 0, batch_size) + md.update_for_spec_dec() + md.use_spec_decoding = False + else: + md.update_for_spec_dec() + self._m3_next_new_tokens = next_new + return draft_tokens + + +def _build_speculation(shell, model_config: ModelConfig): + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise ValueError( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + _set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return draft_config, draft, worker, use_separate + + +# -------------------------------------------------------------------------------------------------------------------- +# Runtime hooks and the shell +# -------------------------------------------------------------------------------------------------------------------- + + +def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: + """The step-loop hooks this target opts into. + + * ``pack_step_inputs``: a decode step's small input copies go out as one upload and one scatter at the graph + replay instead of a chain of copies the next graph waits on; the scatter is parallel (a program per copy), issued + by the decode graph's own prologue (this target's one-launch C++ op, ``m3_step_prologue``) from a window copied + on a side stream while the previous graph runs, so no operation is left between the step graphs and the replay + does not start with a host copy (its copy-engine -> compute turnaround was 10.3 us at c1). Sampled speculative + batches stage their Philox seed/offset uploads, and the accepted-draft counts, into the same pack; the sampler's + slot ids stay in a device buffer re-uploaded only when they change. + * ``skip_unchanged_h2d``: steady decode re-uploads slot indices, query lengths and page tables that rarely change. + * ``kv_lens_update_once``: a speculative overlap step re-derives the MSA KV-length metadata once, after the + forward's KV-length correction (the first pass is rewritten in full by the second). + * ``spec_offsets_op``: the overlap scheduler's speculative step correction (position_ids += offsets, kv_lens += + offsets) as this target's one plain launch (``m3_spec_offsets``) instead of two ATen adds at the head of the + decode graph: the same int32 adds. + * ``loop_on_execution_stream``: the overlap loop's own stream work and the forward on one stream, so step N+1's + inputs wait on graph N in program order rather than through two cross-stream event hops. + * ``merge_tuned_tactics_across_ranks``: tensor-parallel peers capture the same GEMM tactics (rank-local profiling + noise otherwise splits them, and the slowest rank sets every step). + * ``skip_greedy_rng_upload``: all-greedy speculative batches never read the Philox seeds. + * ``spec_step_fold``: the spec sampler's slot stores and the engine's next-step input gathers; the decode forward + stores its outputs and stages the next step's inputs at its own tail (``M3SpecStepFold``). + * ``sampler_d2h_side_stream``: the spec sampler copies its per-slot store to the host on a side stream, overlapping + the next step's input staging. + * ``iteration_timing_one_event``: the iteration log times the device step with one event per boundary on a side + stream (the reported step times are the same). + * MSA: identity selections for short sequences (and short/long graph families), the one-launch KV-length update, + unchanged block-table copies skipped, pool views built once, block-id tables copied from the page-index views + (the stock builder makes a Python list of the whole padded row: on the Vera CPU 282 us per call at one request + and 561 us at two, against 8 and 13 us). + """ + return RuntimeHooks( + first_content_emit_iters=FIRST_CONTENT_EMIT_ITERS, + mxfp8_pinned_tactics=True, + late_compile_detect=True, + pack_step_inputs=True, + pack_parallel_scatter=True, + pack_in_graph=True, + pack_early_copy=True, + pack_prologue_op=m3_step_prologue, + skip_unchanged_h2d=True, + kv_lens_update_once=True, + spec_offsets_op=m3_spec_offsets, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + pack_spec_rng_uploads=True, + spec_small_h2d=True, + spec_step_fold=spec_fold, + sampler_d2h_side_stream=True, + iteration_timing_one_event=True, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=_msa_kv_update_kernel.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + block_ids_from_views=True, + ), + ) + + +class ModelingV2MiniMaxM3Nvfp4Sm107Tp2( + DecoderModelForCausalLM[MiniMaxM3Tp2ModelingV2, PretrainedConfig] +): + """MiniMax-M3 NVFP4 / sm_107 / tp2, text only, with one-model Eagle3 or no speculation. + + The shell owns the vocab-parallel BF16 ``lm_head`` and the logits processor (``DecoderModelForCausalLM``); with + speculation, the drafter (``M3Eagle3Draft``, its weights from the speculative checkpoint through + ``load_draft_weights``) and the worker (``M3Eagle3Worker``) that verifies and drafts each step. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + assert torch.cuda.get_device_capability() == _SM, ( + f"this target is built for sm_{_SM[0]}{_SM[1]}, the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + # Under modeling_v2='require' a covered decode shape that cannot take its fast path raises. + set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + model_config = _text_model_config(model_config) + cfg = model_config.pretrained_config + spec_config = getattr(model_config, "spec_config", None) + if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: + # The one-model Eagle3 layer has no NVFP4 trtllm-gen decode kernel: keep its modules and the shared + # draft cache on FP8 while the target stays NVFP4. + model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 + super().__init__( + MiniMaxM3Tp2ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self.spec_config = spec_config + self.draft_model: Optional[M3Eagle3Draft] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[M3Eagle3Worker] = None + self.use_separate_draft_kv_cache = False + self.layer_idx = FINAL_HIDDEN_LAYER_ID + self._spec_fold: Optional[M3SpecStepFold] = None + if spec_config is not None: + ( + self.draft_config, + self.draft_model, + self.spec_worker, + self.use_separate_draft_kv_cache, + ) = _build_speculation(self, model_config) + self._spec_fold = M3SpecStepFold() + self.runtime_hooks = _runtime_hooks(self._spec_fold) + if mpi_rank() == 0: + logger.info( + f"[M3 sm_107 tp2] routing: modeling_v2 target __tp2, decode rows 1..{WIDE_MAX_TOKENS} (decode layer " + f"1..{DECODE_MAX_TOKENS}; fast-path cells up to 128 rows; 129..{WIDE_MAX_TOKENS} take the stock-exact " + "paths)" + ) + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights: Dict, weight_mapper=None) -> None: + """The drafter's checkpoint (``speculative_model``), loaded by the engine after the target's.""" + _weights.DRAFT_WEIGHTS.load(self.draft_model, weights) + + def set_guided_decoder(self, guided_decoder) -> bool: + """Mirrors ``SpecDecOneEngineForCausalLM.set_guided_decoder``: guided decoding with speculation is installed + on the worker.""" + if hasattr(self.spec_worker, "set_guided_decoder"): + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + def cache_derived_state(self) -> None: + """Build the targets and every decode kernel once the weights are real, before the executor's warmup; with + speculation also the drafter's full-context decode, the fused verify's coded target head, the die-placed + target head copy, the fused draft loop's head copy and the drafter linears' low-M kernels.""" + super().cache_derived_state() + self.model.post_load_weights() + if self.draft_model is not None and DRAFTER_FULLCTX.keys: + path = DRAFTER_FULLCTX.path + recipes = {k: FULLCTX_DECODE_RECIPES[k] for k in DRAFTER_FULLCTX.keys} + installed = m3_drafter_attn.install( + self.draft_model, + recipes, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and m3_fullctx_decode_kernel.has_own_reducer(), + get_env_enable_pdl(), + fused_rope_mode=( + DRAFTER_ROPE_MODE + if DRAFTER_ROPE_FUSED + and hasattr(torch.ops.trtllm, "minimax_m3_sm107tp2_drafter_rope_fp8_kv") + else None + ), + on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), + on_taken=lambda rows: taken(path, rows), + on_uncovered=lambda reason: uncovered(path, reason), + on_missed=lambda rows, reason: missed(path, rows, reason), + ) + logger.info(f"[M3 prebuild] {path} installed on {installed} drafter attention layer(s)") + if self.spec_worker is not None: + self.spec_worker.prebuild(self.lm_head, self.draft_model) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + """One engine step: the target forward, then the logits (or, with speculation, the worker's verify and + draft, and on a decode step the spec fold's store of its outputs). Text only.""" + multimodal_params = kwargs.pop("multimodal_params", None) or [] + if any(p is not None and getattr(p, "multimodal_data", None) for p in multimodal_params): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + worker = self.spec_worker + capture = _plain_capture + capture_target = None + if worker is not None: + capture = worker.capture_hidden_states + capture_target = worker.capture_target + if spec_metadata is not None: + worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + capture=capture, + capture_target=capture_target, + **kwargs, + ) + if getattr(attn_metadata, "csa2_remote_tail_mode", None) == "source": + # A context worker stopped at the encoder/decoder boundary: a disposable token for the context-only + # completion, and no draft worker. + return self.logits_processor.forward(hidden_states, self.lm_head, attn_metadata, False) + if spec_metadata is not None and spec_metadata.is_layer_capture(FINAL_HIDDEN_LAYER_ID): + capture(spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + if worker is None: + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], self.lm_head, attn_metadata, True + ) + spec_input_ids = input_ids if input_ids is not None else kwargs.get("orig_input_ids") + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if spec_input_ids is not None: + spec_input_ids = spec_input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens) + outputs = worker( + input_ids=spec_input_ids, + position_ids=spec_position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + spec_fold = self._spec_fold + if ( + spec_fold is not None + and spec_fold.in_forward + and attn_metadata.num_contexts == 0 + and isinstance(outputs, dict) + ): + # The decode forward stores its outputs into the spec sampler's store and stages the next step's inputs + # at its tail (one launch inside the step graph). + spec_fold.store_in_forward(outputs) + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/__init__.py new file mode 100644 index 000000000000..b8c7426e3933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Activation ops of this target: the gated MLPs' fused SwiGLU.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..daea97df9d12 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SiluAndMul(OpWrapper): + """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with + ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain + SwiGLU ``silu(gate) * up``. + + The MXFP8 dense MLPs (layers 0-2, intermediate 6144 per rank) and the MoE's shared expert (intermediate 1536 per + rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 9216 per + rank) runs it plain. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=1536, oai=True)), + Cell(why="shared expert, an 8-row decode verify", spec=dict(m=8, inter=1536, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=6144, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=6144, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=9216, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=9216, oai=False), + ), + ) + + note = """ + The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two + differ only by the fast sigmoid's last bits (the test's band). + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + gate, up = x.float().chunk(2, dim=-1) + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + a = 1.0 if alpha is None else alpha + b = 0.0 if beta is None else beta + return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" + + +m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/__init__.py new file mode 100644 index 000000000000..a7f7fee7a234 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention ops of this target: MSA block selection, the sparse decode and the KV-length update.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..789ae79db88e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per +request count (modeling's ``DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 16 +KV_HEADS = 1 +HEAD_DIM = 128 +#: The split this target hands trtllm-gen at one request (modeling's DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model +#: maximum the in-tree step passes otherwise. +CAPPED = 16384 +MODEL_MAX = 1 << 20 + + +def dense_attention_reference( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first + ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's + tokens in order). Returns fp32 [total_q, heads, head_dim].""" + batch, kv_heads, max_len, _ = keys.shape + heads = q.shape[1] + group = heads // kv_heads + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + pos = torch.arange(max_len, device=q.device) + with true_fp32_matmul(): + for b in range(batch): + full = int(seq_lens[b]) + rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] + limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) + valid = pos[None, :] < limit[:, None] # [qlen, max_len] + for h in range(kv_heads): + qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] + logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + probs = torch.softmax(logits, dim=-1) + out[ + b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group + ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) + return out + + +def gather_paged_kv( + kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool + (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + page = pool.shape[2] + pages = (int(seq_lens.max()) + page - 1) // page + slots = block_table[:, :pages].long() + k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] + v = pool[slots * subpages + 1] + batch, _, kv_heads, _, d = k.shape + + def order(t): + return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) + + return order(k), order(v) + + +class M3DenseDecode(OpWrapper): + """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len + ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. + + Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into + ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: + every CTA still covers its share of each request's actual length, which may exceed it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="two requests keep the model maximum (no cap at two requests)", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), + ), + ) + + note = """ + The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped + one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV + band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + minimax_m3_trtllm_gen_dense_decode( + q, + kv_cache_manager, + layer_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + max_seq_len=max_seq_len, + max_num_requests=max_num_requests, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). + + ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + if keys.dtype == torch.float8_e4m3fn: + q = fp8_e4m3_saturate(q) + return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> None: + # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. + assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" + # Each verify row attends a causal prefix of at least one key. + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + + +m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..c2b23aa18750 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention on the TRTLLM backend: RoPE, KV-cache append and causal attention.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class M3DraftAttention(OpWrapper): + """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the + backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's + NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` + ``num_tokens``); returns [num_tokens, q heads x 128] bf16. + + ``backend`` (``TrtllmAttention``) and ``metadata`` are the runtime's objects for the drafter layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="first draft pass of a 2-request verify (2 x 4 rows)", spec=dict(reqs=2, qlen=4)), + Cell(why="a later draft pass, one row per request", spec=dict(reqs=2, qlen=1)), + Cell(why="a context chunk's draft rows", spec=dict(reqs=1, context=512)), + ) + + note = """ + `reference` is exact (float64) causal softmax attention: Q rotated by the backend's NeoX RoPE at each token's + position, over the drafter layer's K / V rows read back from its FP8 cache (dequantized by the backend's + ``kv_scale_quant_orig``). What remains between the two is the FMHA's own arithmetic (Q's rounding to the + kernel's input precision, its P rounding and accumulation order), which the test bounds. + """ + + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors grouped KV heads (query head h reads KV head h // (heads / kv heads); two per rank at tp2) and + un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + # [kv heads, tokens, dim]: each KV head's rows in token order across the pages. + keys = ( + pool[pages, 0] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + values = ( + pool[pages, 1] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + keys = keys.repeat_interleave(group, dim=0) # [heads, tokens, dim] + values = values.repeat_interleave(group, dim=0) + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] + rows = q[first : first + n] + x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,hkd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.einsum( + "thk,hkd->thd", torch.softmax(logits, dim=-1), values + ) + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) + + +m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..2fecad0e9afa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q +out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3Fp8IndexerQkNormRope(OpWrapper): + """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and + writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, + then RoPE on its leading ``rotary_dim`` channels (NeoX). + + The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; + the fused layers' producer writes the index cache itself. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, two index heads per rank", spec=dict(m=1, heads=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, heads=2)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=2)), + ) + + note = """ + `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 + before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of + values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. + """ + + def __call__( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + index_k_cache, + out_cache_loc, + num_heads_q, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + position_ids, + ) + + def reference( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, + then the E4M3 cast.""" + out = fused_qk_norm_rope( + qk.clone(), + num_heads_q, + 1, + 0, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) + return ( + q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), + k.to(torch.float8_e4m3fn), + ) + + +m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py new file mode 100644 index 000000000000..f79da95e389c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' full-context decode attention at TP2 on Rubin (``kernels/m3_fullctx_decode``): PrimTS split-KV +decode at the recipe modeling's ``FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) + +from ...kernels import m3_fullctx_decode as kernel +from .m3_dense_decode import dense_attention_reference, gather_paged_kv + +__all__ = ["m3_fullctx_decode"] + +HEADS = kernel.HEADS +KV_HEADS = kernel.KV_HEADS +HEAD_DIM = kernel.HEAD_DIM +#: The plan capacity (the per-request K/V length bound PrimTS sizes its policy for): dispatch's +#: FULLCTX_DECODE_MAX_KV_LEN, which the forward's prebuilt plans use. +MAX_KV_LEN = 1 << 21 + + +class M3FullCtxDecode(OpWrapper): + """Decode attention of every query row over its request's full causal context in the paged FP8 pool, written + into ``output`` (BF16 [T, heads, head_dim]) in place. + + ``recipe`` is ``modeling's FULLCTX_DECODE_RECIPES[(qlen, requests)]``: the PrimTS Q tile and the split-KV fanout. + ``own_reducer`` merges the splits on the BF16 op (modeling's ``FULLCTX_DECODE_OWN_REDUCER``) rather than PrimTS's + FP16 reducer and a cast. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over an AgentX-length (165k) context", + spec=dict(kv_lens=(165525,), qlen=4), + ), + Cell( + why="one request, one draft-pass row, over a 165k context", + spec=dict(kv_lens=(165522,), qlen=1), + ), + Cell( + why="two requests verifying 1 + 3 rows at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165525, 120011), qlen=4), + ), + Cell( + why="two requests, one row each, at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165522, 120008), qlen=1), + ), + Cell( + why="one request verifying 1 + 3 rows over a short (4.1k) context: most splits are empty", + spec=dict(kv_lens=(4104,), qlen=4), + ), + ) + + note = """ + Same operands and arithmetic as the stock trtllm-gen call (E4M3 q, E4M3 K/V, P quantized to E4M3 per split, + fp32 accumulation): the split differs, so the result differs from stock as a stock split change does (the P + quantization is per split). Deterministic, and the same bits eager and in a CUDA graph. Held to the fp32 + reference within the FP8-KV band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + shape = (decode_query_len, int(seq_lens.shape[0])) + kernel.prebuild(q.device, {shape: recipe}, max_kv_len, own_reducer, get_env_enable_pdl()) + plan = kernel.plan_for(q.device, *shape, own_reducer) + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + k_pages, v_pages = kernel.kv_page_views(pool, subpages) + plan.run(fp8_e4m3_saturate(q), k_pages, v_pages, block_table, seq_lens, sm_scale, output) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V, q rounded to E4M3 first as the kernel runs it.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + return dense_attention_reference( + fp8_e4m3_saturate(q), keys, values, seq_lens, decode_query_len, sm_scale + ) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> None: + pool, _ = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + assert pool.dtype == torch.float8_e4m3fn, "an FP8 E4M3 pool" + assert tuple(pool.shape[1:]) == (KV_HEADS, kernel.PAGE, HEAD_DIM), ( + "the TP2 dense-layer pool geometry" + ) + assert q.shape[1:] == (HEADS, HEAD_DIM), "32 query heads of 128 per rank" + assert int(q.shape[0]) == int(seq_lens.shape[0]) * decode_query_len, ( + "q covers exactly the batch" + ) + if torch.cuda.is_current_stream_capturing(): + return # the length checks read the device: eager calls only + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + assert bool((seq_lens <= max_kv_len).all()), "a request longer than the plan capacity" + + +m3_fullctx_decode = M3FullCtxDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..e9ed36eec6f8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``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"] + +#: The wide target's 9..48-request buckets (modeling's INDEX_SCORE_WIDE_SPLITS, repeated here because the ops do not +#: import the dispatch; test_modeling_v2_m3_dispatch checks that every covered request count is certified, so they +#: cannot drift). +_WIDE_SPLITS = ((10, 127), (13, 97), (17, 74), (23, 55), (31, 41), (39, 32), (48, 26)) + + +def _wide_split(requests: int) -> int: + return next(s for hi, s in _WIDE_SPLITS if requests <= hi) + + +def _wide_lens(requests: int) -> tuple: + """Short, mid and edge contexts (one 165521-token request at 12 / 20 / 32 / 46 requests).""" + base = (129, 4000, 20000, 65536, 300, 1000, 8193, 16384) + lens = [base[i % len(base)] for i in range(requests)] + if requests in (12, 20, 32, 46): + lens[0] = 165521 + return tuple(lens) + + +PAGE = 128 +HEAD_DIM = 128 + + +def compiled_scorer(key: tuple): + """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss + (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" + from ...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 kernel at every request count it covers) issues the Q and first K tiles right after the + grid wait from a pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise + every older block of the CTA is prefetched. Neither, nor ``split_k``, changes a score. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over 165521 tokens (AgentX ISL) at TP2: 1536 CTAs, early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=2, split_k=1536, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most of the 1536 CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=2, split_k=1536, early=True), + ), + Cell( + why="one request of 400000 tokens: 3125 blocks, so the 1536 CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=2, split_k=1536, early=True), + ), + Cell( + why="one request decoding one row over 65536 tokens", + spec=dict(seq_lens=(65536,), qlen=1, heads=2, split_k=1536, early=True), + ), + Cell( + why="two requests verifying 4 rows each over 165521 and 20000 tokens: 768 CTAs each, early issue", + spec=dict(seq_lens=(165521, 20000), qlen=4, heads=2, split_k=768, early=True), + ), + Cell( + why="two requests, one of 300 tokens and one of 262144 (2048 blocks: the 768 CTAs loop)", + spec=dict(seq_lens=(300, 262144), qlen=4, heads=2, split_k=768, early=True), + ), + Cell( + why="three requests verifying 4 rows each over 165521, 148968 and 132416 tokens: 424 CTAs each", + spec=dict(seq_lens=(165521, 148968, 132416), qlen=4, heads=2, split_k=424, early=True), + ), + Cell( + why="three requests, 129 / 323584 / 20000 tokens (2528 blocks: the 424 CTAs loop; the short one drains)", + spec=dict(seq_lens=(129, 323584, 20000), qlen=4, heads=2, split_k=424, early=True), + ), + Cell( + why="four requests verifying 4 rows each over 165521, 148968, 132416 and 115864 tokens: 318 CTAs each", + spec=dict( + seq_lens=(165521, 148968, 132416, 115864), qlen=4, heads=2, split_k=318, early=True + ), + ), + Cell( + why="four requests, 300 / 400000 / 129 / 65536 tokens (3125 blocks: the 318 CTAs loop)", + spec=dict(seq_lens=(300, 400000, 129, 65536), qlen=4, heads=2, split_k=318, early=True), + ), + Cell( + why="five requests over 165521 down to 99312 tokens (0.6 x): 254 CTAs each", + spec=dict( + seq_lens=(165521, 148968, 132416, 115864, 99312), + qlen=4, + heads=2, + split_k=254, + early=True, + ), + ), + Cell( + why="six requests, 129 / 400000 / 300 / 20000 / 65536 / 1000: 212 CTAs (3125 blocks loop; short drain)", + spec=dict( + seq_lens=(129, 400000, 300, 20000, 65536, 1000), + qlen=4, + heads=2, + split_k=212, + early=True, + ), + ), + Cell( + why="seven requests over 323584 down to 129414 tokens (0.4 x): 181 CTAs each, all loop", + spec=dict( + seq_lens=(323584, 291225, 258867, 226508, 194150, 161792, 129433), + qlen=4, + heads=2, + split_k=181, + early=True, + ), + ), + Cell( + why="eight requests, 165521 down to 49656 tokens (0.3 x) and one of 129: 159 CTAs each", + spec=dict( + seq_lens=(165521, 148968, 132416, 115864, 99312, 82760, 49656, 129), + qlen=4, + heads=2, + split_k=159, + early=True, + ), + ), + ) + tuple( + Cell( + why=f"{r} requests at {_wide_split(r)} CTAs each (W-IDX1): " + + ("one long request" if r in (12, 20, 32, 46) else "short and mid contexts"), + spec=dict(seq_lens=_wide_lens(r), qlen=4, heads=2, split_k=_wide_split(r), early=True), + ) + for r in range(9, 49) + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..6ec67e73edd4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MiniMax-M3 MSA attention backend's step: KV write, sparse top-k block selection, FMHA.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + +from .m3_sparse_decode_lc import dequant_nvfp4_pages + + +class M3MsaAttention(OpWrapper): + """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` + [T, heads * 128] (bf16, the step's live rows). + + * ``k`` / ``v`` given (the dense layers): written into the layer's paged main cache first (NVFP4 with + ``kv_scale_orig_quant``, else FP8 / BF16 as the pool is). None: a producer already wrote the caches. + * ``sparse``: the indexer scores the layer's index-K cache (already written) with ``idx_q`` and publishes the + top-k blocks; the FMHA attends those blocks. Otherwise the FMHA attends the whole page table. + + The FMHA dispatches on the step's metadata (context rows: the MSA prefill FMHA; generation rows: the MSA decode + FMHA, the dense layers' trtllm-gen decode at their capped KV split). ``backend`` and ``metadata`` are the runtime's + objects for this layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dense layer, a prefill chunk, FP8 pool", + spec=dict(layer="dense", phase="context", tokens=1024), + ), + Cell( + why="dense layer, a decode step of 2 requests", + spec=dict(layer="dense", phase="generation", reqs=2), + ), + Cell( + why="sparse layer, a prefill chunk beyond 16 blocks, NVFP4 pool", + spec=dict(layer="sparse", phase="context", tokens=4096), + ), + Cell( + why="sparse layer, a mixed step (context + generation rows)", + spec=dict(layer="sparse", phase="mixed", tokens=512, reqs=2), + ), + ) + + note = """ + `reference` is exact (float64) softmax attention over the cache rows the step attends, read back from the paged + cache the call wrote (or a producer wrote before it) and dequantized as the FMHA dequantizes them: FP8 rows as + E4M3, NVFP4 rows as E2M1 x E4M3 block scale x the per-tensor ``kv_scale_quant_orig``. A sparse layer's top-k blocks + are re-selected from ``idx_q`` and the index-K cache with the backend's rule (per-block max over the causal + keys, the query's own block forced in, ties to the lower block). On an FP8 pool the FMHA takes Q as a plain E4M3 + cast, which the reference applies too. What remains between the two is the FMHA's own arithmetic (its P + rounding and accumulation order), which the test bounds. + """ + + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors grouped KV heads (query + head h reads KV head h // (heads / kv heads); two per rank at tp2), each KV head selecting its own top-k + blocks from its own index heads against the rank's single index-K head. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + assert config.num_index_heads % kv_heads == 0 + idx_group = config.num_index_heads // kv_heads + + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + qpos = cached[row] + torch.arange(n, device=q.device) + causal = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + for g in range(kv_heads): + if nvfp4: + keys = dequant_nvfp4_pages( + pool[:, 0].view(torch.uint8), + scale_pool[:, 0], + pages, + quad_scales=False, + kv_head=g, + ) + values = dequant_nvfp4_pages( + pool[:, 1].view(torch.uint8), + scale_pool[:, 1], + pages, + quad_scales=True, + kv_head=g, + ) + keys, values = keys * k_global, values * v_global + else: + keys = pool[pages, 0, g].to(torch.float64) + values = pool[pages, 1, g].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + allowed = causal + if sparse: + # One index-K head per rank, shared by every index head (the cache is [pages, 1, page, dim]). + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = _selected_blocks( + idx64[first : first + n, g * idx_group : (g + 1) * idx_group], + idx_keys, + causal, + qpos, + page, + config, + ) + allowed = causal & chosen.repeat_interleave(page, dim=1)[:, :length] + h0, h1 = g * group, (g + 1) * group + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = ( + torch.einsum("thd,kd->thk", q64[first + lo : first + hi, h0:h1], keys) + * sm_scale + ) + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi, h0:h1] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + (ids >= n_valid - config.local_blocks) & (ids < n_valid) + ) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen + + +m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py new file mode 100644 index 000000000000..ff9a05c0233b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the TP2 heads, with its constant loads +staged before the PDL grid-dependency wait and an early dependents trigger +(``kernels/minimaxM3Sm107Tp2/producerPw2.cu``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +__all__ = ["m3_producer"] + +HEAD_DIM = 128 +ROTARY_DIM = 64 +PAGE = 128 +#: Heads per rank at TP2 (MiniMax-M3: 64 query, 4 KV and 4 index heads in all; one index-K head per token). +HEADS_Q = 32 +HEADS_KV = 2 +HEADS_INDEX = 2 +#: Packed heads per token: q, k, v, index-q and the index-k head. +HEADS_PACKED = HEADS_Q + 2 * HEADS_KV + HEADS_INDEX + 1 +#: NVFP4 block: 16 values share one E4M3 scale. +FP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _norm_rope( + x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(FP4_BLOCK) + + +def scale_index(within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return (within_page // 4) * 32 + cols * 4 + within_page % 4 + + +class M3Producer(OpWrapper): + """The NVFP4 QKV + indexer producer of a sparse layer: ``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert`` at the + TP2 heads (32 query, two KV and two index heads), head dim 128 and rotary dim 64, with the same arguments and + outputs bit for bit. + + ``packed`` [T, 39 * 128] BF16 holds each token's 32 q heads, 2 k heads, 2 v heads, 2 index-q heads and the + index-k head. Every head but v gets the Gemma RMSNorm with its weight and NeoX RoPE on its first 64 dims + (``rotary_cos_sin[position]``). Returns (q FP8 [T, 32, 128], index-q FP8 [T, 2, 128]); writes the token's slot + ``out_cache_loc[t]`` (skipped when negative or past the pool): each KV head's k and v as NVFP4 into its block of + the HND pools (16-value blocks, E4M3 scales against ``kv_quant_scale[1]`` / ``[2]``, the V scales in the 4-token + swizzle) and index-k as FP8. ``early`` triggers the dependent kernel on entry (its outputs are read only after the + dependent's own grid wait); ``prewait_step_inputs`` reads the position ids, the slots and the RoPE rows before the + kernel's own grid wait (none of them may be written by a PDL predecessor). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token (c1 without speculation)", spec=dict(tokens=1, scale=1.0)), + Cell(why="a c1 verify: 4 tokens", spec=dict(tokens=4, scale=3.0)), + Cell(why="a c4 verify: 16 tokens, small activations", spec=dict(tokens=16, scale=1e-3)), + Cell( + why="a c8 verify: 32 tokens, large activations, a skipped and an out-of-pool slot", + spec=dict(tokens=32, scale=40.0, bad_slots=True), + ), + Cell(why="the widest decode step: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="a prefill chunk of 1000 tokens, early trigger off", + spec=dict(tokens=1000, scale=1.0, early=False), + ), + ) + + note = """ + The kernel is the native producer's arithmetic, compiled with the library's --use_fast_math, so it is held to that + kernel bit for bit on every output and cache byte. The torch reference states the math in fp32: the FP8 outputs are + held to it within one E4M3 step, and the NVFP4 K/V, decoded, within one E2M1 step of their block scale. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + q, index_q = torch.ops.trtllm.minimax_m3_sm107_tp2_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + prewait_step_inputs, + ) + return q, index_q + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, 32, 128], "index_q" [T, 2, 128], per token and KV head "k" + and "v" [T, 2, 128], and "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl``, + ``early`` and ``prewait_step_inputs`` do not change them.""" + t = packed.shape[0] + x = packed.float().view(t, HEADS_PACKED, HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = rope[:, 0], rope[:, 1] + + def heads(first, count, weight): + return torch.stack( + [_norm_rope(x[:, h], weight, cos, sin, eps) for h in range(first, first + count)], + dim=1, + ) + + k0 = HEADS_Q + v0 = k0 + HEADS_KV + iq0 = v0 + HEADS_KV + return { + "q": heads(0, HEADS_Q, q_weight), + "k": heads(k0, HEADS_KV, k_weight), + "v": x[:, v0 : v0 + HEADS_KV], + "index_q": heads(iq0, HEADS_INDEX, index_q_weight), + "index_k": _norm_rope(x[:, iq0 + HEADS_INDEX], index_k_weight, cos, sin, eps), + } + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> None: + assert heads_q == HEADS_Q, f"the producer is built for {HEADS_Q} query heads, not {heads_q}" + t = packed.shape[0] + # The RoPE row is read at the position with no bound. + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + # Two tokens writing one slot race. + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" + + +m3_producer = M3Producer() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..e253c70a53aa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3QkNormRopeFp8(OpWrapper): + """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded + once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place + entry. + + The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps + of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp2 heads (32 q, 2 kv)", spec=dict(m=1, q=32, kv=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, q=32, kv=2)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=32, kv=2)), + ) + + note = """ + `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from + fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. + """ + + def __call__( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + + def reference( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" + out = fused_qk_norm_rope( + qkv.clone(), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return out.to(torch.float8_e4m3fn) + + +m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..21d361cac90c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into +the paged caches.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3QkvInsert(OpWrapper): + """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the + partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and + index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. + + ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool + [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 + [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp2 + heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), + Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), + ) + + note = """ + The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` + returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + common = ( + num_heads, + num_kv_heads, + num_index_heads, + head_dim, + rotary_dim, + eps, + q_norm, + k_norm, + index_q_norm, + index_k_norm, + rotary_cos_sin, + position_ids, + ) + if kv_scale is not None: + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common + ) + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, index_k_cache, out_cache_loc, *common + ) + + def reference( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of + Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as + packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" + m = packed.shape[0] + sizes = [ + num_heads * head_dim, + num_kv_heads * head_dim, + num_kv_heads * head_dim, + num_index_heads * head_dim, + head_dim, + ] + q, k, v, iq, ik = packed.float().split(sizes, dim=-1) + cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) + half = rotary_dim // 2 + + def norm_rope(x, w, heads): + x = x.reshape(m, heads, head_dim) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) + x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] + c, s = cos.unsqueeze(1), sin.unsqueeze(1) + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) + + return { + "q": norm_rope(q, q_norm, num_heads), + "index_q": norm_rope(iq, index_q_norm, num_index_heads), + "k": norm_rope(k, k_norm, num_kv_heads), + "v": v.reshape(m, num_kv_heads, head_dim), + "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], + } + + +m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..76841cca732d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SelectBlocks(OpWrapper): + """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), + with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the + row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` + backs the result head-major so ``permute(1, 0, 2)`` is contiguous. + + The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="eight decode rows of 17..128 blocks", + spec=dict(total_q=8, n_blocks=128, mixed=True), + ), + Cell( + why="one decode row at the 128-block edge", + spec=dict(total_q=1, n_blocks=128, mixed=False), + ), + Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), + ) + + note = """ + Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. + """ + + def __call__( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + return torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output + ) + + def reference( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + """Plain torch selection; ``head_major_output`` changes the backing only.""" + heads, n_blocks, total_q = max_score_kv.shape + out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) + blocks = torch.arange(n_blocks, device=max_score_kv.device) + for t in range(total_q): + nv = int(n_valid_blocks[t]) + for h in range(heads): + s = max_score_kv[h, :, t].clone().double() + s[blocks >= nv] = float("-inf") + forced = (blocks < min(init_blocks, nv)) | ( + (blocks >= nv - local_blocks) & (blocks < nv) + ) + s[forced] = float("inf") + order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] + ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) + out[t, h, : ids.numel()] = ids + return out + + +m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..66f7f33e8452 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). +MAX_BLOCKS = 9216 +#: Blocks the selection keeps per (query, KV head). +TOPK = 16 +#: Scores the forcing substitutes: initial blocks rank above local ones, and both above any finite score. +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 + +# M3's checkpoint at TP2 (two KV heads per rank): no forced initial blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=2, init_blocks=0, local_blocks=1) + + +class M3SelectLong(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query's top-16 blocks. + + ``scores`` is the per-KV-head block maxima [num_kv_heads, num_blocks, + total_q] the CuTe DSL index scorer wrote (amax over the KV head's index + heads). Per (query, KV head) row the op keeps the ``n_valid_blocks[query]`` + leading blocks (clamped to ``[0, num_blocks]``), forces the last + ``local_blocks`` of them (counted from the *unclamped* valid count) to + ``LOCAL_SCORE`` and the first ``init_blocks`` to ``INIT_SCORE`` -- local wins + where both apply -- and selects the 16 largest, ties to the smaller block. + A pick whose forced score is -inf is emitted as -1. The ids are written + ascending with a -1 tail. + + Same contract as ``torch.ops.trtllm.minimax_m3_select_blocks`` (the in-tree + ``select_blocks_from_maxscore``) and bit-identical to it, except on a row + where more than 2048 valid blocks tie at the 16th-largest score: there the + in-tree op fills the tied slots in atomic order (its staging holds 2048 + candidates), while this op breaks the tie to the smallest blocks as the + contract says. Up to 2048 tied blocks both take the smallest. The test checks + both against ``reference``. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the TP2 engine's 5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + note = """ + Certified at two KV heads per rank (TP2, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`). The kernel takes any head count and + forcing; other forcing is not driven here. + + `scores` may be strided (the op reads it by stride); the target hands over + the scorer's contiguous view, which is what the cells drive. Rows of at + most 128 blocks are within the kernel's range, but the target keeps those on + the in-tree op, so they are not certified here. + + Ordering is by the descending-score key `orderedKey`: a positive NaN ranks + above +inf, +0 above -0, and a negative NaN below -inf -- which, unlike + -inf, is emitted as a block id rather than -1. The cells drive positive + NaN only. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sm107_tp2_select_long( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and + ``trigger_phase`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() # [total_q, num_blocks] + forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where( + in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max + ) + k = min(TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) + ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> None: + """The inputs this op takes and answers wrongly. + + Negative forcing counts are accepted and push the forced ranges past the + row instead of failing; the in-tree op this one stands in for refuses + them, so a caller relying on this one would silently diverge from it. + An ``out`` overlapping ``scores`` is overwritten while it is still read. + """ + assert init_blocks >= 0 and local_blocks >= 0, ( + f"init_blocks ({init_blocks}) and local_blocks ({local_blocks}) must be non-negative; the " + "op accepts negative counts and selects as if the forced ranges were elsewhere" + ) + storage = scores.untyped_storage() + assert not (storage.data_ptr() <= out.data_ptr() < storage.data_ptr() + storage.nbytes()), ( + "out must not alias scores: the op writes it while scores are still read" + ) + + +m3_select_long = M3SelectLong() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py new file mode 100644 index 000000000000..191b4e3132ae --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_select_long import _M3, MAX_BLOCKS, M3SelectLong + + +class M3SelectLongDsl(OpWrapper): + """``m3_select_long``'s contract, cells and reference, run by ``kernels/m3_select_long_dsl``: the same selection + bit for bit as the 1024-thread ``minimax_m3_sm107_tp2_select_long`` (also on rows where more than 2048 blocks tie + at the 16th-largest score), from one 384-thread CTA per row. Returns ``out``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the TP2 engine's 5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + """``m3_select_long``'s reference: the selection is the same bit for bit.""" + return M3SelectLong.reference( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> None: + """``m3_select_long``'s domain.""" + M3SelectLong.is_valid( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + + note = """ + Certified at two KV heads per rank (TP2, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`), and bit for bit against the + 1024-thread `m3_select_long` on every cell. The kernel takes any head count + and forcing; other forcing is not driven here. + + The kernel is compiled per (PDL launch, early trigger, trigger phase) with + symbolic shapes and strides; the target compiles its keys before warmup + (`kernels/m3_select_long_dsl.prebuild`), and a call outside them compiles + on first use. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + from ...kernels.m3_select_long_dsl import select_long_dsl + + if scores.shape[1] > MAX_BLOCKS: + raise RuntimeError( + f"m3_select_long_dsl supports at most {MAX_BLOCKS} blocks, got {scores.shape[1]}" + ) + select_long_dsl( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + +m3_select_long_dsl = M3SelectLongDsl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..0a8613019283 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This target's block-sparse decode attention (Triton, with the one-launch long-context route).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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 (two KV heads) with its block and per-tensor + scales, and BF16 ``output`` [T, 32, 128]. Returns ``output``. + + Which kernel runs is decided inside, per step: a selected (non-identity) table on the TP2 geometry goes to the + one-launch ``m3_sparse_decode_lc`` kernel while modeling's ``SPARSE_DECODE_LC`` covers it; an identity table + (every request within 16 blocks, staged by the metadata), or a selected table with that range killed, runs the + Triton per-row producer, then the LSE merge. (The multi-query producer needs one KV head: never at TP2.) Every + route also publishes the MXFP8 o_proj hand-off to ``m3_prequant.SLOT``, bit for bit + ``torch.ops.trtllm.mxfp8_quantize(output, True)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="identity selection, decode without speculation at batch 8: the Triton producer + merge", + spec=dict( + seq_lens=(1, 128, 129, 700, 2048, 1500, 64, 0), decode_query_len=1, identity=True + ), + ), + Cell( + why="identity selection, Eagle3 verify at one request (T=4): the Triton producer + merge", + spec=dict(seq_lens=(1900,), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 2 (T=8), a padding request", + spec=dict(seq_lens=(2048, 0), decode_query_len=4, identity=True), + ), + Cell( + why="selected blocks, Eagle3 verify at one request (T=4): the one-launch long-context kernel", + spec=dict(seq_lens=(40000,), decode_query_len=4, identity=False), + ), + Cell( + why="selected blocks, Eagle3 verify at batch 2 (T=8): the one-launch long-context kernel", + spec=dict(seq_lens=(9000, 40000), decode_query_len=4, identity=False), + ), + ) + + note = """ + Certified only as the target calls it: two KV heads of 16 query heads each, head_dim 128, FP8 q, packed NVFP4 + cache, BF16 output, `num_topk_chunks` left to the dispatcher. The function also takes BF16 / FP8 caches, other + head geometries and an explicit split-K factor (the in-tree tests drive those); none is certified here. + + The MXFP8 hand-off of the long-context route is certified on its own entry (`m3_sparse_decode_lc`); the Triton + merge's (`_store_mxfp8_head`) is held to `torch.ops.trtllm.mxfp8_quantize` by this entry's test. + """ + + def __call__( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + kv_scale=kv_scale, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + num_topk_chunks=num_topk_chunks, + identity_blocks=identity_blocks, + ) + return output + + def reference( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + """``output`` as the call is supposed to write it: float64 block-sparse attention over the + dequantized NVFP4 blocks (``sparse_attention_reference``), rounded once to ``output``'s dtype. + + Only the packed NVFP4 cache is mirrored -- the target's only configuration -- one KV head at a time: KV head + ``h`` serves query heads ``[h * group, (h + 1) * group)`` with its own top-k row. ``kv_scale``, + ``num_topk_chunks`` and ``identity_blocks`` do not change what is computed. + """ + assert k_block_scale is not None, "mirrors the NVFP4 cache only" + kv_heads = k_paged.shape[1] + group = q.shape[1] // kv_heads + heads = [] + for h in range(kv_heads): + heads.append( + sparse_attention_reference( + q[:, h * group : (h + 1) * group], + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[h], + block_table, + seq_lens, + decode_query_len, + sm_scale, + kv_head=h, + ) + ) + return torch.cat(heads, dim=1).to(output.dtype) + + def is_valid( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) + + +m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py new file mode 100644 index 000000000000..2733eb1e7524 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py @@ -0,0 +1,474 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch long-context sparse decode over the selected NVFP4 blocks (16-CTA clusters), SM107 TP2 geometry. + +Also the plain-torch block-sparse attention and MXFP8 hand-off mirrors the +Triton sparse decode entry (``m3_sparse_decode``) reuses: both compute the same +thing over the same cache. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The TP2 geometry the kernel is compiled for: 16 query heads on each of 2 KV heads per rank. +HEADS_PER_KV = 16 +KV_HEADS = 2 +NUM_HEADS = HEADS_PER_KV * KV_HEADS +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +#: NVFP4 elements per E4M3 block scale. +SF_VEC = 16 +#: MXFP8 elements per UE8M0 scale, and the o_proj width the hand-off is laid out for. +MX_BLOCK = 32 +WIDTH = NUM_HEADS * HEAD_DIM + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, + block_scale: torch.Tensor, + pages: torch.Tensor, + quad_scales: bool, + kv_head: int = 0, +) -> torch.Tensor: + """Dequantize KV head ``kv_head`` of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, kv_head].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, kv_head].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, + kv_head: int = 0, +) -> torch.Tensor: + """Block-sparse decode attention over KV head ``kv_head`` of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages( + k, k_block_scale, pages, quad_scales=False, kv_head=kv_head + ).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True, kv_head=kv_head).reshape( + -1, HEAD_DIM + ) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, + sf [num_rows_alloc * WIDTH // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) + num_sf_k_tiles = WIDTH // 128 + sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +class M3SparseDecodeLc(OpWrapper): + """Sparse decode attention, split-K merge and MXFP8 hand-off of one TP2 decode step in one launch. + + One cluster of 16 CTAs per (query token, KV head): CTA ``r`` attends the + token's ``r``-th selected block of that head (``topk[h, tok, r]``), and the + cluster merges the 16 partials over distributed shared memory into the + head's 16 BF16 rows of ``out[tok]``. With ``q8`` / ``sf`` it also writes + the row's MXFP8 copy and swizzled UE8M0 scales for the o_proj, bitwise + ``torch.ops.trtllm.mxfp8_quantize(out, True)``. All operands are the byte + views ``m3_attn_lc.m3_lc_attn`` hands over (E4M3 ``q``, packed E2M1 pages, + E4M3 block scales). Returns ``out``. + + Per block: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16 + (exact), QK and PV on FP16 tensor cores with FP32 accumulation, the softmax + probabilities rounded to FP16 for PV, log2-domain softmax with + ``ex2.approx``. ``linear_softmax`` picks the denominator recurrence (the two + agree to rounding); SM107 keeps the log-domain recurrence of the stock + Triton producer. ``prewait`` stages the request's block-table row in shared + memory before the grid dependency wait; it does not change a value. + ``n1_ids`` (int32 [2, rows, 16], dense rows, initialized to -1) holds the + layer's selected block ids one decode step back: each CTA prefetches its + slot's block into L2 before the grid dependency wait and overwrites the slot + with this step's id (-1 past the token's block count); rows at or past + ``rows`` skip both. A hint: it does not change a value. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="Eagle3 verify at one request, 165k context: T=4, log-domain recurrence, MXFP8 hand-off", + spec=dict( + seq_lens=(165521,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="Eagle3 verify at two requests, long and medium context: T=8, no hand-off", + spec=dict( + seq_lens=(165521, 98765), + decode_query_len=4, + linear_softmax=False, + mxfp8=False, + prewait=False, + ), + ), + Cell( + why="Eagle3 verify at two requests: T=8 with the MXFP8 hand-off", + spec=dict( + seq_lens=(165521, 12345), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="two short requests with -1 tails and a zero-length padding request: T=8", + spec=dict( + seq_lens=(130, 0), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="decode without speculation at two requests: T=2", + spec=dict( + seq_lens=(165521, 5000), + decode_query_len=1, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="the wide target's largest hand-off step: 32 requests, T=128 (one full scale tile), mixed lengths " + "with tile edges and a zero-length padding request, MXFP8 hand-off", + spec=dict( + seq_lens=tuple((165521, 20000, 1281, 129)[i % 4] for i in range(31)) + (0,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="the linear denominator and pre-wait staging variants, driven at T=8", + spec=dict( + seq_lens=(165521, 20000), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + prewait=True, + ), + ), + ) + + note = """ + Certified: the eight instantiated variants (`linear_softmax` x `prewait` x + hand-off on / off). The hand-off cells are certified bit-exactly against + `mxfp8_handoff_reference` applied to the op's own BF16 rows, which equals + `torch.ops.trtllm.mxfp8_quantize(out, True)`; the BF16 rows themselves within + the band the test derives. `prewait` and the hand-off do not change the BF16 + rows. + + `topk` is the head-major [2, T, 16] table; it is read for the first + `min(16, ceil(kv_len / 128))` entries of each token; a -1 among them skips + that CTA's block, entries past them are never read. `use_pdl` only changes + the launch attribute. The hand-off assumes `num_sf_k_tiles = 32` (a + 4096-wide o_proj input) and buffers with at least as many rows as the step; + `MERGE_MXFP8` sizes them at 128 rows. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + n1_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sm107_tp2_sparse_decode_lc( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + decode_query_len, + sm_scale, + linear_softmax, + use_pdl, + prewait, + n1_ids, + ) + return out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + n1_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 4096] the op is supposed to write: ``sparse_attention_reference`` per KV head in + float64, rounded once. ``linear_softmax``, ``use_pdl``, ``prewait`` and ``n1_ids`` do not change what is + computed; the hand-off buffers are checked separately with ``mxfp8_handoff_reference``. + """ + q3 = q.view(-1, NUM_HEADS, HEAD_DIM) + attn = torch.cat( + [ + sparse_attention_reference( + q3[:, h * HEADS_PER_KV : (h + 1) * HEADS_PER_KV], + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk[h], + block_table, + seq_lens, + decode_query_len, + sm_scale, + kv_head=h, + ) + for h in range(KV_HEADS) + ], + dim=1, + ) + return attn.reshape(out.shape).to(out.dtype) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + n1_ids: Optional[torch.Tensor] = None, + ) -> None: + """The inputs this op takes and answers wrongly (``m3_attn_lc.lc_applies`` is the target's form of the + layout half). + + The kernel addresses a page's token rows at a fixed 64-byte (data) and + 8-byte (scale) pitch and loads them 16 / 8 bytes at a time, so only the + page and KV-head strides are read from the views. It reads the selected + block ids, the block table and the hand-off buffers unchecked. Reads + ``topk`` and ``seq_lens`` back to the host: a test-only cost. + """ + for name, t in (("k", k), ("v", v)): + assert ( + t.dim() == 4 and t.shape[1] == KV_HEADS and t.shape[2:] == (PAGE, HEAD_DIM // 2) + ), f"{name} must be the [pages, 2, 128, 64] byte view of the rank's KV heads" + assert t.stride(3) == 1 and t.stride(2) == HEAD_DIM // 2, ( + f"{name} rows must be dense 64-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0 and t.stride(1) % 16 == 0, ( + f"{name} pages and heads must be 16-byte aligned" + ) + assert v.stride() == k.stride(), "k and v share their strides" + for name, t in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + assert ( + t.dim() == 4 + and t.shape[1] == KV_HEADS + and t.stride(3) == 1 + and t.stride(2) == HEAD_DIM // SF_VEC + ), f"{name} rows must be dense 8-byte rows: the kernel ignores the row stride" + assert t.data_ptr() % 8 == 0 and t.stride(0) % 8 == 0 and t.stride(1) % 8 == 0, ( + f"{name} pages and heads must be 8-byte aligned" + ) + assert q.dim() == 3 and q.shape[1:] == (NUM_HEADS, HEAD_DIM), "q must be [T, 32, 128]" + assert q.data_ptr() % 16 == 0 and q.stride(0) % 16 == 0 and q.stride(1) % 16 == 0, ( + "q rows are loaded 16 bytes at a time and must be 16-byte aligned" + ) + assert k_global_scale.numel() == 1 and v_global_scale.numel() == 1, ( + "global scales are single values" + ) + total_q = q.shape[0] + assert topk.dim() == 3 and topk.shape[0] == KV_HEADS and topk.shape[1] == total_q, ( + "topk is the head-major [2, T, 16] table" + ) + assert ( + total_q % decode_query_len == 0 and seq_lens.numel() >= total_q // decode_query_len + ), ( + "every query token must belong to a request: seq_lens is read at tok // decode_query_len unchecked" + ) + seq = seq_lens.tolist() + num_cols = block_table.shape[1] + for h in range(KV_HEADS): + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk[h, tok, : min(TOPK, nblocks)] + assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( + f"head {h} token {tok}: selected block ids must be -1 or below the request's block count; " + "the block table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + f"num_sf_k_tiles must be {WIDTH // 128}: the hand-off rows are {WIDTH} wide and a smaller " + "value overlaps them silently" + ) + assert ( + q8.numel() >= total_q * WIDTH + and sf is not None + and sf.numel() >= ((total_q + 127) // 128 * 128 * (WIDTH // MX_BLOCK)) + ), "the hand-off buffers must hold every row of the step: they are written unchecked" + + +m3_sparse_decode_lc = M3SparseDecodeLc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..43166b472eca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The MSA pool's page: one sparse block is one 128-token KV page. +TOKENS_PER_BLOCK = 128 + + +class MsaKvUpdate(OpWrapper): + """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. + + ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its + own buffers at the head of every captured speculative decode step. One + single-CTA launch writes, for ``batch`` requests and ``total_q`` query + tokens (all int32): + + * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, + ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); + * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` + attends up to, through ``block_table``; + * ``n_valid_blocks[t]``, when given, the block count top-k selection is + bounded by. + + It returns None; the cells read the four buffers. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation at one request: qo = 1, every buffer written", + spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="decode without speculation at the captured batch 8", + spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", + spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), + ), + Cell( + why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", + spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), + ), + Cell( + why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", + spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), + ), + Cell( + why="positions crossing a page boundary: the slot moves to the previous page of the request", + spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), + ), + ) + + note = """ + The kernel covers any batch up to the block size's loop (it strides over + 256 threads), and any page size; only `tokens_per_block = 128`, the MSA + pool's page, is certified. + + The op does not bounds-check the request index a token reads + (`q_batch_row`) or the block-table column it derives: both come from the + metadata's own staging, which keeps them in range by construction, and an + out-of-range value reads whatever memory sits there. `is_valid` states both. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.minimax_m3_msa_kv_update( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + + def reference( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], + n_valid_blocks[:total_q])`` as the op is supposed to write them. + + The in-tree metadata's own torch re-derivation + (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook + call), statement for statement, returning new tensors rather than + writing the buffers. The buffer arguments only say which outputs exist: + a ``None`` buffer gives a ``None`` result. + """ + page = int(tokens_per_block) + kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) + cu_kv = None + if cu_kv_lens is not None: + cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) + cu_kv[1:] = torch.cumsum(kv_true, 0) + qbr = q_batch_row[:total_q].to(torch.long) + qo_dev = seq_lens[:batch] + pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(block_table.shape[1]) - 1 + ) + slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + n_valid = None + if n_valid_blocks is not None: + n_valid = torch.div( + (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" + ).to(torch.int32) + return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid + + def is_valid( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + """The inputs this op takes and answers wrongly. + + Every operand but the block table is read as a dense run from its first + element, so a strided view reads the wrong elements. The request index of + each token and the table row it selects are not bounds-checked. Reads + ``q_batch_row`` back to the host: a test-only cost. + """ + for name, t in ( + ("kv_lens", kv_lens), + ("kv_lens_staged", kv_lens_staged), + ("seq_lens", seq_lens), + ("q_batch_row", q_batch_row), + ("q_intra", q_intra), + ("msa_seq_lens", msa_seq_lens), + ("out_cache_loc", out_cache_loc), + ("cu_kv_lens", cu_kv_lens), + ("n_valid_blocks", n_valid_blocks), + ): + assert t is None or t.is_contiguous(), ( + f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" + ) + rows = q_batch_row[:total_q] + assert total_q == 0 or ( + int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) + ), ( + "q_batch_row must index a request below batch (and a block-table row): the op reads the " + "request's lengths and table row unchecked" + ) + + +msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/__init__.py new file mode 100644 index 000000000000..5da86fe24bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective-communication entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..7f480855c170 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams + + +class M3AllReduce(OpWrapper): + """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the + plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: + ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; + returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the + communicator's per call. + + The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps + beyond the multicast kernel's rows. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="tp2, a prefill chunk", spec=dict(world=2, m=2048, hidden=6144)), + Cell(why="tp2, an 8-row generation step", spec=dict(world=2, m=8, hidden=6144)), + Cell( + why="tp2, plain sum of eight drafter rows", + spec=dict(world=2, m=8, hidden=6144, plain=True), + ), + ) + + note = """ + The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then + adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum + to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. + """ + + def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): + return comm(x, all_reduce_params=params) + + def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): + """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" + total = torch.stack([x.float() for x in xs]).sum(0) + if params is None: + return total.to(xs[0].dtype) + updated = (params.residual.float() + total).to(xs[0].dtype) + h = updated.float() + normed = ( + h + * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) + * params.norm_weight.float() + ) + return normed.to(xs[0].dtype), updated + + +m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..883f38e0803d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,319 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode layer boundary, 2 ranks on SM107.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer + fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; + only the stored residual is rounded to BF16. + + The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are + summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic + and the normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + pre = total.to(torch.bfloat16).float() + residual.float() + updated = pre.to(torch.bfloat16) + u = pre.double() if drafter else updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +class M3McAllreduceNorm(OpWrapper): + """One TP group's boundary: ``updated = bf16(residual + bf16(sum_r bf16(x_r + x2_r)))``, ``normed = + RMSNorm(updated)``: bitwise the stock all-reduce + fused residual RMSNorm (+ mxfp8_quantize) of SM107 TP2. + + Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls + its own view ``ws_uc`` (re-arming the words it read) and rotates the per-CTA call parities in ``flags``; the + workspace is the ``McAllReduceWorkspace`` bucket of ``maxt`` tokens. All ranks must make the same calls in the + same stream order. ``prepushed`` 2 means a producer already pushed ``x`` (and launched this grid after its + pushes); 3 that two producers pushed ``x`` into exchange set A and ``x2`` into set B. Returns + ``[normed, updated]`` or, with ``quant``, also the MXFP8 copy of ``normed`` (E4M3 rows, swizzled UE8M0). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The tp2 decode boundaries (hidden 6144, 2 ranks, eps 1e-6): every decode step fits the 8-token bucket, epilogue + # option 7 at 1, 2, 7, 8 tokens and 15 at 3..6 (kernels/m3_mc_allreduce.AR_OPTS_BY_M; every option is + # bit-identical). + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token at the attention boundary, x only, option 7", + spec=dict(m=1, maxt=8, opts=7, x2=False, quant=False, prepushed=0), + ), + Cell( + why="MoE boundary at 4 rows (one request verifying 1 + 3): shared + routed summed in-kernel, MXFP8 copy", + spec=dict(m=4, maxt=8, opts=15, x2=True, quant=True, prepushed=0), + ), + Cell( + why="MoE boundary at 8 rows (two requests), MXFP8 copy for the next QKV", + spec=dict(m=8, maxt=8, opts=7, x2=True, quant=True, prepushed=0), + ), + Cell( + why="attention boundary whose 8 rows the o_proj GEMV already pushed (prepushed 2)", + spec=dict(m=8, maxt=8, opts=7, x2=False, quant=False, prepushed=2), + ), + # Every other decode row count (modeling's MC_ALLREDUCE), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} rows: a boundary row count between the ones above", + spec=dict( + m=m, + maxt=8, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + x2=m % 2 == 0, + quant=m % 3 == 0, + prepushed=0, + ), + ) + for m in range(1, 9) + if m not in (1, 4, 8) + ), + # B3, the batch-4 decode steps (modeling's MC_ALLREDUCE past 8 rows): every row count 9..16 on the 16-token + # bucket (option 15), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} rows on the 16-token bucket (batch-4 decode step)", + spec=dict(m=m, maxt=16, opts=15, x2=m % 2 == 0, quant=m % 3 == 0, prepushed=0), + ) + for m in range(9, 17) + ), + # The batch-8 decode steps: every row count 17..32 on the 48-token bucket (option 15), the second addend and + # the MXFP8 copy alternating. + *( + Cell( + why=f"{m} rows on the 48-token bucket (batch-8 decode step)", + spec=dict(m=m, maxt=48, opts=15, x2=m % 2 == 0, quant=m % 3 == 0, prepushed=0), + ) + for m in range(17, 33) + ), + # The Eagle3 drafter boundaries (modeling's DRAFT_MC_AR, SpecTail): option 16 (the drafter's stock rounding) on + # 23 / 31 at every draft-pass row count 1..8 (pass 0: 4 x requests, later passes: requests), x only. + *( + Cell( + why=f"Eagle3 drafter boundary at {m} rows (opts {23 if m in (1, 2, 7, 8) else 31})", + spec=dict( + m=m, + maxt=8, + opts=23 if m in (1, 2, 7, 8) else 31, + x2=False, + quant=False, + prepushed=0, + drafter=True, + ), + ) + for m in range(1, 9) + ), + # Every other row count the o_proj pushes (modeling's ATTN_PUSH). + *( + Cell( + why=f"{m} rows prepushed by the o_proj GEMV", + spec=dict( + m=m, + maxt=8, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + x2=False, + quant=False, + prepushed=2, + ), + ) + for m in range(1, 8) + ), + # B6, the batch-4 decode steps (modeling's ATTN_PUSH past 8 rows): every row count 9..16 the o_proj pushes into + # the 16-token bucket (option 15). + *( + Cell( + why=f"{m} rows prepushed by the o_proj GEMV on the 16-token bucket (batch-4 decode step)", + spec=dict(m=m, maxt=16, opts=15, x2=False, quant=False, prepushed=2), + ) + for m in range(9, 17) + ), + ) + + note = """ + The epilogue option bits (1 early residual / weight load, 2 st.async cluster exchange, 4 early parity rotation, + 8 max carveout) are bit-identical by construction; only the four (bucket, option) pairs of the sm_103 kernel are + built. SM107 TP2 against the stock AUTO all-reduce + fused residual RMSNorm + mxfp8_quantize, rows 1..24, x alone + and x + x2: normed, updated and the MXFP8 copy (data and the scale bytes of valid rows) bitwise equal. `prepushed` + 1 (pushed, grid wait kept) is accepted by the op but no caller passes it. `pdl` changes only launch overlap and is + driven off in the cells. + + `normed` is gated to one BF16 rounding step: the kernel's RMS statistic is a warp / cluster tree of FP32 sums + and its `rsqrtf` is approximate, so a row whose scaled value lands next to a BF16 midpoint can round either + way. `updated` and the MXFP8 copy of the op's own `normed` are bit-exact. + """ + + def __call__( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_sm107_tp2_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + None, + None, + maxt, + quant, + prepushed, + opts, + cap_out, + cap_w, + cap_eps, + ) + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant, drafter) + + def is_valid( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> None: + h = residual.shape[-1] + # The op checks norm_w's element count, not its layout: a strided view is read as packed memory. + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + # The workspace is addressed as [set][parity][maxt][world][H / 2] words and never bounds-checked: a smaller + # buffer (another bucket's, or one built for fewer ranks) is polled and re-armed out of range. + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + # Flags are per (token, cluster CTA) of the bucket; the op checks the count, not that they belong to the + # workspace. Flags of another bucket hold another call history, and the parity they select is not the one the + # peers push into: the poll then waits on the other half forever or reads last call's rows. + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + + +m3_mc_allreduce_norm = M3McAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..9861ed2f0513 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize + multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode MoE boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_mc_allreduce_norm import boundary + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class M3McFinalizeAllreduceNorm(OpWrapper): + """The MoE boundary with the routed experts' finalize folded in. + + Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, + ``scales`` fp32 or bf16), then the group's all-reduce, residual add and RMSNorm exactly as + ``m3_mc_allreduce_norm``: same workspace, flags and call-order contract. Returns ``[normed, updated]`` or, with + ``quant``, also the MXFP8 copy of ``normed``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Not taken at tp2 (modeling's MC_FINALIZE is empty: the decode step is at most 8 rows); the cells are the tp4 + # target's, on the SM107 kernel. The tp4 target takes this path from 9 tokens (top-4 of 128 experts, fp32 routing + # weights, the shared expert always present): the 16- and 48-token buckets, epilogue option 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="the smallest finalize call (9 tokens, 16-token bucket) with the next QKV's MXFP8 copy", + spec=dict(m=9, maxt=16, opts=15, topk=4, quant=True, dropped=0), + ), + Cell( + why="a full 16-token bucket with dropped expert slots", + spec=dict(m=16, maxt=16, opts=15, topk=4, quant=False, dropped=5), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=24, maxt=48, opts=15, topk=4, quant=False, dropped=0), + ), + Cell( + why="the 48-token bucket full, MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, topk=4, quant=True, dropped=3), + ), + # Every other row count of the unfinalized MoE (MC_FINALIZE): the bucket and options the target picks. + *( + Cell( + why=f"{m} tokens: a finalize row count between the ones above", + spec=dict( + m=m, + maxt=16 if m <= 16 else 48, + opts=15, + topk=4, + quant=m % 2 == 0, + dropped=m % 4, + ), + ) + for m in range(9, 49) + if m not in (9, 16, 24, 48) + ), + ) + + note = """ + BF16 routing weights and the 8-token bucket are accepted by the op (option 7 there) but the target only enters + this path at 9 tokens and up, with the built-in MoE's fp32 weights; neither is a cell. The rest is as in + `m3_mc_allreduce_norm`'s note. + """ + + def __call__( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_sm107_tp2_mc_allreduce_norm( + fc2, + shared, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + 0, + opts, + cap_out, + cap_w, + cap_eps, + ) + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> None: + h = residual.shape[-1] + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + # The kernel gathers fc2[perm] without a bound: an index past the FC2 rows reads whatever follows them. + lo, hi = int(perm.min()), int(perm.max()) + assert lo >= -1 and hi < fc2.shape[0], ( + f"perm spans [{lo}, {hi}]; entries must be -1 (dropped) or a row of fc2 (< {fc2.shape[0]})" + ) + + +m3_mc_finalize_allreduce_norm = M3McFinalizeAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/__init__.py new file mode 100644 index 000000000000..e9f4d8e20b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedding ops of this target: the generic token-embedding gather.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..71fce6b253cb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="an 8-row decode verify, int64 ids", + spec=dict(rows=8, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ) + + note = """ + An exact gather: `reference` must match bit for bit. + """ + + def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) + + def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] + return rows * valid.unsqueeze(-1).to(table.dtype) + + +m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/__init__.py new file mode 100644 index 000000000000..9cc374bdf1cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMM entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..3ed3cc97a229 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3Bf16Linear(OpWrapper): + """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. + + The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), + o_proj (this rank's input columns), gate_up and down projections run it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=4608, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=4096)), + Cell(why="drafter gate_up, one row", spec=dict(m=1, n=18432, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=9216)), + ) + + note = """ + cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x, w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return (x.float() @ w.float().t()).to(x.dtype) + + +m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py new file mode 100644 index 000000000000..f1b619a13564 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense MLP's swigluoai activation as one C++ launch with PDL and an early release (DENSE_SWIGLU_CPP).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# The dense layers' swigluoai parameters, which the kernel hard-codes (modeling_minimaxm3._build_swiglu_oai_dense_mlp). +LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 + + +class M3DenseSwiglu(OpWrapper): + """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, + rounded once to bf16: trtllm::minimax_m3_sm107_tp2_swiglu_oai (minimaxM3Sm107Tp2/swigluOai.cu), bit for bit the + Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the + grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights + during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even + half width; ``__call__`` returns None otherwise or for other activation parameters (the caller runs the Triton + op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP2 dense MLP shard: gate | up 12288 wide (d = 6144), at the decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP2 dense MLP at 1 row", spec=dict(rows=1, d=6144)), + Cell(why="the TP2 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=6144)), + Cell(why="the TP2 dense MLP at 8 rows", spec=dict(rows=8, d=6144)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the + alpha-scaled ex2.approx sigmoid, div.full), checked over every bf16 gate and every bf16 up pattern. + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float] = LIMIT, + alpha: Optional[float] = ALPHA, + beta: Optional[float] = BETA, + use_pdl: bool = True, + ) -> Optional[torch.Tensor]: + if ( + (limit, alpha, beta) != (LIMIT, ALPHA, BETA) + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 4 != 0 + or x.data_ptr() % 4 != 0 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp2_swiglu_oai(x, out, bool(use_pdl)) + return out + + def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_dense_swiglu = M3DenseSwiglu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py new file mode 100644 index 000000000000..427ee2598f91 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 x BF16 -> BF16 low-M GEMM of the Eagle3 drafter's linears at fixed FlashInfer tactics (SM10x).""" + +from __future__ import annotations + +from typing import Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: (N, K) of the TP2 drafter shards -> ((max rows, kind, tactic), ...), rows ascending. Kind "d" is FlashInfer's direct +#: SIMT kernel with tactic (block, outputs per block, rows per block); kind "s" is its split-K tensor-core kernel with +#: tactic (mma_m, mma_n, split_k, ab_stages). A row count is served by the first entry whose max rows holds it; the +#: rows modeling's DRAFT_LOWM covers are only the ones measured (1, and 4 where a tactic beats cuBLAS at no larger error +#: vs fp64). Rubin SM107, mclk locked, 4 rotating weight sets, CUDA-graph timing, us cuBLAS -> low-M at 1 / 4 rows: +#: qkv 13.18 -> 9.71 / 13.31 -> 11.72, o 7.17 -> 5.33 / 7.12 -> 5.78, gate_up 22.27 -> 17.70 / -, down 12.52 -> 9.76 / +#: 12.59 -> 10.13, fc 21.50 -> 17.77 / -. +TACTICS: Mapping[Tuple[int, int], tuple] = { + # qkv + (4608, 12288): ((1, "d", (256, 2, 1)), (4, "d", (64, 4, 4))), + # o_proj + (6144, 4096): ( + (1, "d", (128, 2, 1)), + (4, "s", (64, 8, 2, 10)), + (8, "s", (64, 8, 2, 10)), + (16, "s", (64, 16, 2, 10)), + ), + # gate_up + (18432, 6144): ((1, "d", (256, 4, 1)),), + # down + (6144, 9216): ( + (1, "d", (128, 2, 1)), + (4, "s", (64, 8, 2, 10)), + (16, "s", (64, 16, 2, 10)), + (32, "s", (64, 32, 2, 6)), + ), + # the Eagle3 fc over the three captured hidden states + (6144, 18432): ((1, "d", (384, 2, 1)), (16, "s", (64, 16, 2, 10))), +} +#: Rows each (N, K) is certified at (a subset of what its table would serve). +ROWS: Mapping[Tuple[int, int], Tuple[int, ...]] = { + (4608, 12288): (1, 4), + (6144, 4096): (1, 4, 8, 16), + (18432, 6144): (1,), + (6144, 9216): (1, 4, 16, 32), + (6144, 18432): (1, 16), +} + + +def tactic_for(n: int, k: int, m: int) -> Optional[Tuple[str, tuple]]: + """(kind, tactic) of an [m, k] x [n, k]^T call, or None when the table does not serve it.""" + for rows, kind, tactic in TACTICS.get((n, k), ()): + if m <= rows: + return kind, tactic + return None + + +class M3DraftLowM(OpWrapper): + """``out`` [M, N] BF16 = ``x`` [M, K] BF16 @ ``weight`` [N, K]^T BF16, FP32 accumulation. + + ``kind`` / ``tactic`` pick FlashInfer's kernel: "d" runs ``run_direct_dense`` (SIMT, one dot product per output in + K order), "s" runs ``run_splitk_dense`` (tcgen05 MMA, K split over ``split_k`` CTAs and summed in FP32). Both + round once to BF16. The kernel is compiled on its first call per (M, N, K, tactic, PDL); the target compiles every + covered call before warmup (``M3Eagle3Draft.prebuild_linears``). + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter qkv, 1 rows: d tactic", spec=dict(nk=(4608, 12288), m=1)), + Cell(why="drafter qkv, 4 rows: d tactic", spec=dict(nk=(4608, 12288), m=4)), + Cell(why="drafter o_proj, 1 rows: d tactic", spec=dict(nk=(6144, 4096), m=1)), + Cell(why="drafter o_proj, 4 rows: s tactic", spec=dict(nk=(6144, 4096), m=4)), + Cell(why="drafter gate_up, 1 rows: d tactic", spec=dict(nk=(18432, 6144), m=1)), + Cell(why="drafter down, 1 rows: d tactic", spec=dict(nk=(6144, 9216), m=1)), + Cell(why="drafter down, 4 rows: s tactic", spec=dict(nk=(6144, 9216), m=4)), + Cell(why="drafter fc, 1 rows: d tactic", spec=dict(nk=(6144, 18432), m=1)), + ) + + note = """ + The table is the TP2 drafter's: its shapes are the per-rank shards of the Eagle3 layer at 2-way tensor + parallelism (64 query and 4 KV heads, intermediate 18432) and the replicated fc, at the b1 decode rows (1 row + in draft passes 1-2, 4 in pass 0). Certified at those shapes and rows only; FlashInfer's kernels take other + shapes, which are not measured here. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + if kind == "d" and ar is not None: + # The boundary push (modeling's DRAFT_AR_PUSH_DIRECT): this target's copy of the direct kernel, whose + # output is FlashInfer's, also stores the rows into the boundary all-reduce's exchange. + from ...kernels.m3_lowm_direct import DirectTactic, run_direct_dense_push + + run_direct_dense_push(x, weight.t(), out, use_pdl, DirectTactic(*tactic), ar) + elif kind == "d": + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + DirectTactic, + run_direct_dense, + ) + + run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) + elif own_splitk: + # This target's copy of the split-K kernel (kernels/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # it also pushes the rows into a drafter boundary all-reduce (modeling's DRAFT_AR_PUSH). + from ...kernels.m3_lowm_splitk import ( + SplitKTactic, + run_splitk_dense, + run_splitk_dense_push, + ) + + if ar is not None: + run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) + else: + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + else: + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + return out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """bf16(x @ weight^T) summed in fp64 and rounded once; ``out``, ``kind``, ``tactic`` and ``use_pdl`` only + choose where and how the op writes.""" + return (x.double() @ weight.double().t()).to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> None: + """The kernels index ``x`` and ``out`` densely: a strided view is read as if it were packed.""" + assert x.dim() == 2 and x.is_contiguous(), "x must be a contiguous [M, K] matrix" + assert weight.dim() == 2 and weight.is_contiguous(), ( + "weight must be a contiguous [N, K] matrix" + ) + assert out.is_contiguous() and tuple(out.shape) == (x.shape[0], weight.shape[0]), ( + "out must be a contiguous [M, N] matrix" + ) + assert ar is None or ( + len(ar) == 5 and ((own_splitk and kind == "s") or (kind == "d" and tactic[1] % 2 == 0)) + ), "ar (the boundary all-reduce push) runs only on this target's split-K or direct copy" + + +m3_draft_lowm = M3DraftLowM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..9702cc0778a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class M3Fp32Linear(OpWrapper): + """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the + fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies + (the engine leaves torch's default).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), + Cell( + why="a generation step above the tensor-core router's 48 rows", + spec=dict(m=64, e=128, k=6144), + ), + ) + + note = """ + Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against + the true fp32 reference. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x.to(torch.float32), w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return x.float() @ w.t() + + +m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..6f04ed1ba6d4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 x MXFP8 GEMV (tcgen05, block-scaled, CuTe DSL) on an activation its producer already quantized.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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", "served_max_rows"] + +# (N, K) of this target's MXFP8 projections at tp2, the GEMV's tuned shapes. +QKV_SPARSE = (4992, 6144) # the fused QKV + index projection (index values disabled) +QKV = (4608, 6144) +O_PROJ = (6144, 4096) +SHARED_GATE_UP = (3072, 6144) +SHARED_DOWN = (6144, 1536) +DENSE_GATE_UP = (12288, 6144) +DENSE_DOWN = (6144, 6144) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "shared gate_up": SHARED_GATE_UP, + "shared down": SHARED_DOWN, + "dense gate_up": DENSE_GATE_UP, + "dense down": DENSE_DOWN, +} + + +def served_cells(named: tuple[Cell, ...], x_bf16: bool, max_m: int) -> tuple[Cell, ...]: + """A cell for every (projection, rows) the GEMV serves up to ``max_m`` rows that ``named`` does not list. + + Certification is per projection: a projection's split-K, ring and MMA width differ, so each (N, K, M) the + decode dispatch can send here is its own cell. + """ + listed = {(cell.spec["nk"], cell.spec["m"]) for cell in named} + return tuple( + Cell(why=f"{name}, {m} rows", spec=dict(nk=nk, m=m)) + for name, nk in PROJECTIONS.items() + for m in range(1, max_m + 1) + if (nk, m) not in listed and _gemv.serves(nk[0], nk[1], m, x_bf16) + ) + + +def served_max_rows(nk: tuple[int, int], x_bf16: bool, lo: int, hi: int) -> int: + """The largest ``m`` in ``lo..hi`` such that the GEMV serves an (N, K) projection at every row count ``lo..m``; + ``lo - 1`` when it declines ``lo``. The decode dispatch raises a projection's rows past the batch-1/2 step only + as far as this (the bf16 GEMV's in-kernel quantize declines the K 6144, split 2 projections at MMA width 16).""" + m = lo - 1 + while m < hi and _gemv.serves(nk[0], nk[1], m + 1, x_bf16): + m += 1 + return m + + +def gemv_reference( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """bf16(x @ dequant(weight)^T) for fp32 activation values x [M, K], summed in fp64 and rounded once.""" + return (x.double() @ mxfp8_dequant(weight, weight_scale).double().t()).to(torch.bfloat16) + + +def gemv_is_valid(weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + n, k = weight.shape + # The kernel reads both from data_ptr() as dense buffers at the layout it was compiled for. + assert weight.dtype == torch.float8_e4m3fn and weight.is_contiguous(), ( + "weight must be contiguous e4m3 [N, K]" + ) + assert weight_scale.numel() == n * k // 32, ( + "weight_scale must be the 128x4-swizzled UE8M0 buffer [N * K / 32]" + ) + + +class M3MxGemv(OpWrapper): + """bf16 [M, N] = xq [M, K] (e4m3, 128x4-swizzled UE8M0 scales ``xsf``) x weight [N, K]^T (e4m3, swizzled + ``weight_scale``), M <= 48; None when the shape is not served (the caller falls back). + + ``trigger_early``: release the PDL dependents before the epilogue. ``cfg``: a (split, ring[, contig, nacc, a_tmem, + st_async]) tuning override; default the table's. ``use_pdl``: default the environment's. ``ar`` = (multicast + workspace address, flags address, rank, world, bucket tokens, mode): also push the rows into the next boundary + all-reduce's exchange (then ``op.last_pushed`` is True). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The projections' (N, K) at the token counts each build differs in: the MMA width (8 columns up to the TP2 decode + # step's 8 rows) is one, so one cell per projection pins the shape, and the rest are every served row count. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 4)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 8 rows: two requests verifying 1 + 3", + spec=dict(nk=QKV_SPARSE, m=8), + ), + Cell( + why="dense-attention QKV, 4 rows: one request verifying 1 + 3", spec=dict(nk=QKV, m=4) + ), + Cell(why="o_proj, one token (split 2)", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 8 rows", spec=dict(nk=O_PROJ, m=8)), + Cell(why="shared gate_up, 4 rows (split 4)", spec=dict(nk=SHARED_GATE_UP, m=4)), + Cell(why="shared down, one token (split 2, K 1536)", spec=dict(nk=SHARED_DOWN, m=1)), + Cell(why="dense gate_up, 8 rows (N 12288)", spec=dict(nk=DENSE_GATE_UP, m=8)), + Cell(why="dense down, 4 rows (K 6144)", spec=dict(nk=DENSE_DOWN, m=4)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling's MX_GEMV_ROWS). + CELLS = CELLS + served_cells(CELLS, x_bf16=False, max_m=48) + + note = """ + The kernel sums in fp32 in split-K order where `reference` sums in fp64: the bf16 outputs agree within the band the + test derives (SM107: at most as many misrounded elements as the stock FlashInfer CuTe DSL GEMM on the same + operands, rand / wide-scale / cancellation / tie / zero inputs). Not covered: the all-reduce push (`ar`), which + needs a multicast exchange; `cfg` overrides beyond the table (the target passes none). Shapes outside the table + and M > 48 return None; the target sends at most 8 rows. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv(xq, xsf, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The dequantized product; the launch options are ignored.""" + return gemv_reference(mxfp8_dequant(xq, xsf), weight, weight_scale) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + m, k = xq.shape + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous(), ( + "xq must be contiguous e4m3 [M, K]" + ) + # A linear-layout [M, K / 32] scale buffer is read as the swizzled one: wrong scales, no error. + assert xsf.numel() >= (m + 127) // 128 * 128 * (k // 32), ( + "xsf must be the 128x4-swizzled scale buffer (rows padded to 128), not the linear one" + ) + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv = M3MxGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..5788c113d449 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 GEMV (CuTe DSL) on a bf16 activation, quantized to MXFP8 in the kernel's prologue.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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, + SHARED_DOWN, + SHARED_GATE_UP, + gemv_is_valid, + gemv_reference, + served_cells, +) + +__all__ = ["m3_mx_gemv_bf16"] + + +class M3MxGemvBf16(OpWrapper): + """bf16 [M, N] = mxfp8_quantize(x) x weight [N, K]^T for a contiguous, 16-byte aligned bf16 ``x`` [M, K]; the + quantization is mxfp8_quantize's (UE8M0 scale per 32 = the smallest power of two >= amax / 448). None when the + shape is not served. The other arguments as ``m3_mx_gemv``'s. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The prologue quantizes one 32-element block per thread and round: the projections at the token counts the + # target sends here (no producer-quantized activation), every TP2 decode row count. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token", spec=dict(nk=QKV_SPARSE, m=1)), + Cell(why="dense-attention QKV, 8 rows", spec=dict(nk=QKV, m=8)), + Cell(why="o_proj, 4 rows (the attention output in bf16)", spec=dict(nk=O_PROJ, m=4)), + Cell(why="shared gate_up, 8 rows", spec=dict(nk=SHARED_GATE_UP, m=8)), + Cell(why="shared down, 4 rows (K 1536)", spec=dict(nk=SHARED_DOWN, m=4)), + Cell( + why="dense gate_up, 8 rows: 24 K tiles per split-2 CTA", + spec=dict(nk=DENSE_GATE_UP, m=8), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling's MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=32) + + note = """ + As `m3_mx_gemv`. The in-kernel quantization is mxfp8_quantize's bit for bit except where the kernel's + rcp.approx(448) and an exact 1/448 put a block's amax on different sides of a power of two (not observed on the + cells). Not covered: `ar`, `cfg` overrides. Returns None for rows beyond 8 rounds of the prologue per CTA. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv_bf16(x, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The MXFP8 round trip of x times the dequantized weight; the launch options are ignored.""" + return gemv_reference(mxfp8_quant_dequant(x), weight, weight_scale) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv_bf16 = M3MxGemvBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..3bfcbc625016 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned + +from .m3_mxfp8_linear import mxfp8_reference + + +class M3Mxfp8Gemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight + [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key + and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The decode dispatch hands it the boundary / merge MXFP8 activation on rows the GEMV does not serve: the + # unsplit shared down projection at 17..32 rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=1536)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=1536)), + ) + + note = """ + Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + """The dequantized product; ``tune`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..9eadb1e0d20a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned + +_BLOCK = 32 + + +def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" + pad_r = (rows + 127) // 128 * 128 + pad_c = (cols + 3) // 4 * 4 + s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) + return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] + + +def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, k = x.shape + s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) + factor = torch.pow(2.0, (s - 127).to(torch.float32)) + return ( + x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) + ).view(rows, k) + + +def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: + """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" + with true_fp32_matmul(): + return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=4992, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=4096, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=6144, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=3072, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=4096, path="tuned")), + ) + + note = """ + Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` + quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and + multiplies dequantized operands in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + if tuned_quantize_gemm: + return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + if flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, output_dtype + ) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are + ignored.""" + m, k = x.shape + blocks = x.float().view(m, k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(-1, keepdim=True) + exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) + scale = torch.pow(2.0, exp) + q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale + with true_fp32_matmul(): + out = q.view(m, k) @ _dequant(weight, weight_scale).t() + return (out * alpha.float()).to(output_dtype) + + +m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py new file mode 100644 index 000000000000..f8deb858c843 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, +and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). + +The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values +``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 +marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a +per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: +per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 +code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved +it decodes to the BF16 weight bit for bit. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_target_head"] + +#: Verify rows per call (requests x (draft length + 1)). +MAX_ROWS = 8 +#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). +CTA_WARPS = 16 +#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. +TARGET_SHARD = 50016 +HIDDEN = 6144 +#: Codes per window: 0..14 decode, 15 marks an exception. +WINDOW = 15 +#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. +MAX_BASE = 254 - WINDOW +#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). +ROWS_PER_CHUNK = 4096 + + +def _pad_rows(weight: torch.Tensor) -> torch.Tensor: + n = weight.shape[0] + tiles = (n + 15) // 16 + if tiles * 16 == n: + return weight + return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) + + +def _fragment_view(w: torch.Tensor) -> torch.Tensor: + """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" + n, k = w.shape + # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) + return ( + w.view(n // 16, 2, 8, k // 16, 2, 4, 2) + .permute(0, 3, 2, 5, 4, 1, 6) + .reshape(n // 16, k // 16, 32, 8) + ) + + +def _exponents(weight: torch.Tensor) -> torch.Tensor: + return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF + + +def _window_base(hist: torch.Tensor) -> int: + """The base of the 15-exponent window holding the most weights (the lowest such base).""" + window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) + return int(window.argmax()) + + +def encode(weight: torch.Tensor, base: Optional[int] = None): + """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's + own window. + + Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile + << 16 | k), exc_w fp32 [E, 2] ({w, 0})). + """ + w = _pad_rows(weight) + n, k = w.shape + tiles, ks = n // 16, k // 16 + assert ks % 4 == 0, "K must be a multiple of 64" + bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF + expo = (bits >> 7) & 0xFF + if base is None: + base = _window_base(torch.bincount(expo.flatten(), minlength=256)) + code = expo - base + exc = (code < 0) | (code >= WINDOW) + code = torch.where(exc, torch.full_like(code, WINDOW), code) + sm = ((bits >> 8) & 0x80) | (bits & 0x7F) + sm = torch.where(exc, torch.zeros_like(sm), sm) + # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. + rows, cols = exc.nonzero(as_tuple=True) + order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) + rows, cols = rows[order], cols[order] + exc_w = torch.stack( + [ + w[rows, cols].float(), + torch.zeros((rows.numel(),), device=w.device), + ], + dim=1, + ) + exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) + exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) + exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( + torch.int32 + ) + # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. + smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) + cf = ( + _fragment_view(code.to(torch.uint8).view(n, k)) + .view(tiles, ks // 4, 4, 32, 8) + .to(torch.int32) + ) + # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. + cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) + for p in range(4): + cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) + cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 + # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. + seg01 = ( + smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) + ) # [.., 2, 32, 16] + seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) + packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] + return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() + + +def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): + """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" + n, _ = weight.shape + assert rows_per_chunk % 16 == 0 + hist = torch.zeros(256, dtype=torch.int64, device=weight.device) + for r in range(0, n, rows_per_chunk): + hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) + base = _window_base(hist) + packed, offs, poss, ws = [], [], [], [] + n_exc = 0 + for r in range(0, n, rows_per_chunk): + p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) + packed.append(p) + offs.append(off[:-1] + n_exc) + poss.append(pos) + ws.append(ew) + n_exc += pos.numel() + exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) + return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) + + +def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: + """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" + tiles, ks = (n + 15) // 16, k // 16 + p = packed.view(tiles, ks // 4, 3, 32, 16) + sm = ( + p[:, :, :2] + .permute(0, 1, 3, 2, 4) + .reshape(tiles, ks // 4, 32, 4, 8) + .permute(0, 1, 3, 2, 4) + .to(torch.int32) + ) + cw = ( + p[:, :, 2] + .contiguous() + .view(torch.int32) + .view(tiles, ks // 4, 32, 4) + .permute(0, 1, 3, 2) + .long() + & 0xFFFFFFFF + ) + codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( + torch.int32 + ) + bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) + bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) + frag = bits.reshape(tiles, ks, 32, 8) + # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). + w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) + return w.to(torch.int16).view(torch.bfloat16)[:n] + + +def decode( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + n: int, + k: int, + row0: int = 0, + row1: Optional[int] = None, +) -> torch.Tensor: + """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). + ``row0`` must be a multiple of 16.""" + row1 = n if row1 is None else row1 + assert row0 % 16 == 0 + t0, t1 = row0 // 16, (row1 + 15) // 16 + unit_bytes = (k // 64) * 3 * 32 * 16 + dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() + counts = (exc_off[1:] - exc_off[:-1]).long() + tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) + rows = tile_of * 16 + (exc_pos >> 16).long() + cols = (exc_pos & 0xFFFF).long() + sel = (rows >= row0) & (rows < row1) + dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) + return dec + + +def round_trip_ok( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + weight: torch.Tensor, + rows_per_chunk: int = ROWS_PER_CHUNK, +) -> bool: + """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent + window below the BF16 infinity exponent.""" + n, k = weight.shape + if not 0 <= base <= MAX_BASE or k % 64: + return False + if ( + packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 + or exc_off.numel() != (n + 15) // 16 + 1 + ): + return False + for r0 in range(0, n, rows_per_chunk): + r1 = min(n, r0 + rows_per_chunk) + dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) + if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): + return False + return True + + +class M3TargetHead(OpWrapper): + """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit + the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in + another summation order). + + ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` + (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are + the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and + # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with + # a shard whose rows are not a multiple of 16. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{m} verify row(s) of the TP4 target head shard", + spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), + ) + for m in range(1, MAX_ROWS + 1) + ), + Cell( + why="4 rows with outliers and exact zeros in many tiles (the exception lists)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), + ), + Cell( + why="2 rows over a shard of 4000 rows (250 tiles), K 1024", + spec=dict(m=2, n=4000, k=1024, exceptions=True), + ), + ) + + note = """ + The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception + adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs + from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); + integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp + order: the result is deterministic. + """ + + def __call__( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_target_head( + w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl + ) + return out + + def reference( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. + + ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" + weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + return logits.to(torch.bfloat16) + + def is_valid( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> None: + # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count + # stores a partial sum. + assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + assert out.data_ptr() != x.data_ptr(), "out must not alias x" + + +m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py new file mode 100644 index 000000000000..335effba7fee --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The target lm_head shard GEMM of a verify step on a two-die GPU, each die reading only its own memory. + +A copy of the head shard [N, K] (BF16) is placed once as one virtual range whose rows [0, S) live in locality domain 0 +and rows [S, N) in domain 1 (:meth:`LocalityMap.place_split`; S on a 2 MiB granule). A covered verify computes +``F.linear`` of each row half of the copy on a green-context stream holding that domain's SMs, and concatenates: each +logit is the in-tree head GEMM's own logit (the same cuBLAS GEMM per output column) at the row counts the cells +certify. The module's own weight is never re-placed: a kernel reading die-placed memory from all SMs runs 1.5-1.7x +slower (SpecTail 769262), so every other reader (uncovered row counts, prefill logits, warmup, fallbacks) keeps the +ordinary weight, and only this op reads the placed copy. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger + +__all__ = ["UgpuHead", "m3_ugpu_head", "prepare_head"] + +TARGET_SHARD = 100032 # the TP2 target head shard rows (vocabulary 200064 / 2) +HIDDEN = 6144 +_ENGAGED: set = set() + + +@dataclass +class UgpuHead: + """A placed head: ``weight`` the die-placed copy (read only by :data:`m3_ugpu_head`), ``source`` the module's own + weight it copies and ``version`` that weight's version counter at the copy (a later in-place change makes the copy + stale: see :func:`current`), ``split`` the first domain-1 row, ``streams`` the two domains' green-context + streams.""" + + weight: torch.Tensor + source: torch.Tensor + version: int + split: int + streams: List[torch.cuda.ExternalStream] + + +def current(head: Optional["UgpuHead"], weight: torch.Tensor) -> bool: + """Whether ``head`` is a placed copy of ``weight`` as it is now (same tensor, not modified since the copy).""" + return head is not None and head.source is weight and weight._version == head.version + + +def _split_row(n: int, row_bytes: int, granule: int) -> Optional[int]: + """The granule-aligned row nearest the middle (the halves' bytes differ by < one granule), or None.""" + import math + + step = granule // math.gcd(granule, row_bytes) # rows per granule-aligned boundary + s = round(n / 2 / step) * step + return s if 0 < s < n else None + + +def prepare_head(lm_head: torch.nn.Module) -> Optional[UgpuHead]: + """A die-placed copy of ``lm_head.weight`` for :data:`m3_ugpu_head` (the module is not changed), or None when the + head does not take the plain ``F.linear`` path this op reproduces, or the device has no two locality domains. + Before capture.""" + from tensorrt_llm._torch.modules import low_m_gemm + + from ...kernels.m3_locality import LocalityMap + + w = getattr(lm_head, "weight", None) + if ( + w is None + or w.dtype != torch.bfloat16 + or w.dim() != 2 + or not w.is_contiguous() + or getattr(lm_head, "bias", None) is not None + or getattr(lm_head, "use_cute_dsl_bf16_gemm", False) + or getattr(lm_head, "use_custom_cublas_mm", False) + or low_m_gemm.LOW_M_GEMM_ACTIVE + or type(getattr(lm_head, "quant_method", None)).__name__ != "UnquantizedLinearMethod" + ): + return None + try: + lm = LocalityMap.get(w.device) + except (ImportError, RuntimeError) as err: + logger.warning(f"[M3 sm_107 tp2] uGPU head: no locality map ({err}); the head stays as is") + return None + s = _split_row(w.shape[0], w.shape[1] * w.element_size(), lm.granularity) + if s is None: + return None + with torch.no_grad(): + placed = lm.place_split(w.detach(), s) + return UgpuHead( + weight=placed, source=w, version=w._version, split=s, streams=lm.domain_streams() + ) + + +class M3UgpuHead(OpWrapper): + """``logits`` [M, N] (BF16) = x [M, K] (BF16) times the placed head W [N, K] transposed, as + ``cat(F.linear(x, W[:S]), F.linear(x, W[S:]))`` with the halves on the two domains' green-context streams + (forked from and joined back into the caller's stream).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The verify rows of one to four requests at draft length 3 (the TP2 GEN batch is 1: every E2E step is 4 rows). + CELLS: tuple[Cell, ...] = ( + Cell( + why="4 verify rows: one request, draft length 3 (the b1 point)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN), + ), + Cell( + why="8 verify rows: two requests, draft length 3", + spec=dict(m=8, n=TARGET_SHARD, k=HIDDEN), + ), + Cell(why="12 verify rows: three requests", spec=dict(m=12, n=TARGET_SHARD, k=HIDDEN)), + Cell( + why="16 verify rows: four requests (the c4 point)", + spec=dict(m=16, n=TARGET_SHARD, k=HIDDEN), + ), + ) + + note = """ + Each output column is computed by cuBLAS over the same weight row and the same input as the in-tree head GEMM; at + the certified row counts the half-N call picks the full call's kernel, so the logits are the in-tree head's bit + for bit (the cells check it on SM107). At 1 and 2 rows cuBLAS picks a different kernel for the half shard: not + covered. Placement only moves pages between the dies; every other reader of the weight is unchanged. + """ + + def __call__(self, x: torch.Tensor, head: UgpuHead) -> torch.Tensor: + key = (x.device, x.shape[0]) + if key not in _ENGAGED: + _ENGAGED.add(key) + logger.info( + f"[M3 sm_107 tp2] fast path engaged: m3_ugpu_head ({x.shape[0]} rows, split {head.split} of " + f"{head.weight.shape[0]}, on {x.device})" + ) + main = torch.cuda.current_stream(x.device) + fork = torch.cuda.Event() + fork.record(main) + parts = [] + for stream, w in zip(head.streams, (head.weight[: head.split], head.weight[head.split :])): + stream.wait_event(fork) + with torch.cuda.stream(stream): + parts.append(F.linear(x, w)) + for stream in head.streams: + done = torch.cuda.Event() + done.record(stream) + main.wait_event(done) + for p in parts: + p.record_stream(main) + return torch.cat(parts, dim=1) + + def reference(self, x: torch.Tensor, head: UgpuHead) -> torch.Tensor: + """The in-tree head GEMM over the whole (placed) weight.""" + return F.linear(x, head.weight) + + +m3_ugpu_head = M3UgpuHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/__init__.py new file mode 100644 index 000000000000..77d7b84d9bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..562a9f392cef --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 1..8 decode tokens in one CuTe DSL launch (``m3_moe``, the sm_107 TP2 build: routed +and shared expert width 1536 per rank): the router, the routed NVFP4 experts and the shared MXFP8 expert, on +caller-owned state (an ``M3MoeState`` and one ``M3MoeLayer`` per MoE layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ...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) + + +_ENGAGED: set = set() + + +class M3MoeEngine(OpWrapper): + """One MiniMax-M3 MoE layer for M (1..8) decode tokens: ``layer.state.run(layer, hidden, out, arp)``. + + ``hidden`` bf16 [M, 6144] contiguous, M = the state's token count. ``layer`` (from ``M3MoeState.layer``) carries the + TRTLLM-Gen NVFP4 expert buffers read in place, the per-expert scalars, the router weight / bias, the activation's + NVFP4 global scale and, on a ``shared`` state, the shared expert's MXFP8 weights; its state carries the compiled + kernel and the workspace every layer on it shares. Returns bf16 [M, 6144]: bf16(shared + routed) on a ``shared`` + state, else the routed partial. ``out``: write there instead of a new tensor. ``arp`` (``ar_push_desc``): also + push the rows into a multicast all-reduce exchange. + + The routed part is TRTLLM-Gen's NVFP4 MoE on the fp32 router logits: sigmoid (tanh form) + correction bias top-4 + of 128 (ties to the lower id), weights sigmoid / sum * 2.0 rounded to bf16; the activation quantized to NVFP4 with + the layer's global scale; per expert the raw FC1 sums, SwiGLU-OAI (alpha / beta / clamp), the NVFP4 requant with + fc31_scale_c, FC2, bf16(fc2_alpha x sum); the fma chain over the top-4 in order. The shared expert is the target's + MXFP8 SwiGLU-OAI MLP. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # One cell per token count: m_max is a trace-time constant, so each is its own build, and the target builds all + # eight (decode batches of 1..8 requests, spec-decode verify rows up to 8). The whys name what the build differs in. + CELLS: tuple[Cell, ...] = ( + Cell( + why="1 token: FC1 -> FC2 records replicated 32x, logits 32x, last phase without trailing FC2", + spec=dict(tokens=1), + ), + Cell( + why="2 tokens: records 32x, logits 32x, the one build without AR_PRE", + spec=dict(tokens=2), + ), + Cell(why="3 tokens: records 8x, logits 32x, 4 FC2-B builder warps", spec=dict(tokens=3)), + Cell( + why="4 tokens: one request verifying 3 drafts (the TP2 b1 point): 6-stage ring, B2IMG hand-off, " + "the 4 two-block FC2 CTAs stream no FC1", + spec=dict(tokens=4), + ), + Cell( + why="5 tokens: single records, a routing warp takes two tokens, logits 32x", + spec=dict(tokens=5), + ), + Cell( + why="6 tokens: single records, routing warps 8.. take tokens 4..5", spec=dict(tokens=6) + ), + Cell(why="7 tokens: the B2IMG global FC1 -> FC2 hand-off", spec=dict(tokens=7)), + Cell( + why="8 tokens: two requests verifying 3 drafts (the TP2 b2 point): 5-stage ring, x read from global " + "(no XSTAGE), no early group 0, the 4 two-block FC2 CTAs stream no FC1", + spec=dict(tokens=8), + ), + ) + + note = """ + The kernel's elementwise math uses the PTX approximations (rcp.approx, ex2.approx, tanh) where `reference` uses + torch's correctly rounded ones, and its GEMMs sum in tensor-core order where `reference` sums in fp64: a + quantization boundary (NVFP4 activation / intermediate code, routing near-tie) can fall on the other side, so the + two agree within the band the test derives, not bit for bit. The routed part is bit for bit the stock TRTLLM-Gen + NVFP4 MoE (fp4_quantize + fp4_block_scale_moe_runner) on the same routing; the test checks that too. + + The all-reduce push (`arp`) needs a multicast exchange: its 2-rank certification (pushed boundary bit for bit the + unpushed one) is comm/_m3_moe_push_op_matrix.py. Not covered: the non-shared state (`shared=False`), which this + target never builds. Every layer on one state + must run in one stream order: the workspace is shared and tagged per call, never reset. + + No `is_valid`: the compiled call checks every buffer's dtype, shape and alignment against its signature and + raises on a mismatch (a misaligned `hidden`, an `out` that is not bf16 [M, 6144], a malformed `arp`). + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + state = layer.state + if (state.device, state.m) not in _ENGAGED: + # The engagement line (once per device and token count, on every rank: the first call, eager or captured). + _ENGAGED.add((state.device, state.m)) + logger.info( + f"[M3 sm_107 tp2] MoE engine m3_moe engaged: {state.m} tokens, {state.num_ctas} routed + " + f"{state.mod.NSH} shared CTAs, {state.mod.STAGES}-stage ring, on {state.device}" + ) + return state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Plain torch over the layer's buffers (see the class docstring); `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=layer.state.shared) + + +m3_moe_engine = M3MoeEngine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py new file mode 100644 index 000000000000..636dfcd0b614 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV of the decode MoE: ``logits[M, E] = float(x) @ W^T`` at one token.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) +#: The kernel builds each block's output scale through two ``rcp.approx.ftz.f32`` reciprocals (~2^-23 relative): +#: a value this close to an e2m1 midpoint may round to either neighbour, everything further out must match. +TIE_WINDOW = 2.0**-20 + + +def nvfp4_quantize( + x: torch.Tensor, sf_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], + the pre-rounding scaled values [M, K] that ``near_tie`` reads). + + Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. + """ + gs = sf_scale.float().reshape(()) + m, k = x.shape + blocks = x.float().view(m, k // 16, 16) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +def near_tie(scaled: torch.Tensor) -> torch.Tensor: + """Elements whose e2m1 rounding the kernel's approximate reciprocals decide (see ``TIE_WINDOW``).""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=scaled.device, dtype=torch.float32) + a = scaled.abs().unsqueeze(-1) + return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) + + +def fp32_logits(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """``float(x) @ w^T`` as a true fp32 product (TF32 off).""" + with true_fp32_matmul(): + return x.float() @ w.t() + + +class M3RouterGemv(OpWrapper): + """fp32 router logits ``float(x[M, K]) @ w[E, K]^T`` with exact fp32 FMA, one CTA per expert. + + ``x`` is BF16 with 1..16 rows, unit column stride and 16-byte rows; ``w`` is the fp32 router weight; K a + multiple of 1024. ``rows`` experts per CTA; only 1 is built. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target takes the GEMV at exactly one token (2..48 go to the tensor-core router): 128 experts, hidden 6144. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token, 128 experts x hidden 6144", + spec=dict(m=1, e=128, k=6144), + ), + ) + + note = """ + The op accepts up to 16 rows; the target sends one, so wider calls are not cells. The logits differ from a true + fp32 product only in accumulation order (per-thread FMA chains over K / 1024 strided column groups, a warp tree + and an 8-warp sum). + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_router_gemv(x, w, rows, pdl) + + def reference(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + """True fp32 product; ``rows`` and ``pdl`` do not enter the arithmetic.""" + return fp32_logits(x, w) + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> None: + # The activation rows are bulk-copied from x's base address; the op checks the row stride but not the base, + # and a misaligned bulk copy faults the context rather than raising. + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + + +m3_router_gemv = M3RouterGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py new file mode 100644 index 000000000000..0306210adaab --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV plus the routed experts' NVFP4 activation, one token.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_router_gemv import fp32_logits, nvfp4_quantize + + +class M3RouterGemvFp4(OpWrapper): + """``[logits fp32 [M, E], x_fp4 [M, K / 2] uint8, x_sf [M * K / 16] uint8]`` from one kernel. + + The logits are ``m3_router_gemv``'s; the NVFP4 activation is ``fp4_quantize(x, sf_scale, 16, False, False)`` + (16-element blocks, E4M3 scales in the linear layout, per-tensor ``sf_scale``) for the routed experts' FC1. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The fused-router decode step at one token: 128 experts, hidden 6144, the experts' input scale. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token: logits for 128 experts and the NVFP4 activation of hidden 6144", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 8.0), + ), + Cell( + why="a global scale that puts block scales near the E4M3 top, where the scale cast saturates", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 2.0), + ), + ) + + note = """ + The e2m1 codes are exact except within `TIE_WINDOW` of a rounding midpoint, where the kernel's approximate + reciprocals decide the direction. Up to 16 rows are accepted; the target sends one. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_gemv_fp4(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """True fp32 logits and the NVFP4 recipe; ``pdl`` does not enter the arithmetic.""" + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return fp32_logits(x, w), packed, sf + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + # A zero or non-finite global scale makes every block scale 0 or NaN and the data codes garbage; the op + # reads the scalar on the device and cannot refuse it. + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_gemv_fp4 = M3RouterGemvFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..9e21113d8483 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TF32 tensor-core router of the decode MoE for 2..48 tokens, optionally with the NVFP4 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from .m3_router_gemv import nvfp4_quantize + + +def tf32_rna(w: torch.Tensor) -> torch.Tensor: + """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" + bits = w.contiguous().view(torch.int32) + rounded = (bits + 0x1000) & ~0x1FFF + finite = (bits & 0x7F800000) != 0x7F800000 + return torch.where(finite, rounded, bits).view(torch.float32) + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target routes 1..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one token, logits only (the wide target's router_path(1) with no fp32 GEMV rows)", + spec=dict(m=1, k=6144, fp4=False), + ), + Cell( + why="one token with the NVFP4 activation (the fused-router path at one row)", + spec=dict(m=1, k=6144, fp4=True), + ), + Cell( + why="two tokens, logits only (the smallest call the target makes here)", + spec=dict(m=2, k=6144, fp4=False), + ), + Cell( + why="8 tokens with the NVFP4 activation: one 8-token MMA tile, every expert tile quantizes one row", + spec=dict(m=8, k=6144, fp4=True), + ), + Cell( + why="13 tokens: a partial second token tile", + spec=dict(m=13, k=6144, fp4=True), + ), + Cell( + why="the 48-token maximum, six token tiles, logits only", + spec=dict(m=48, k=6144, fp4=False), + ), + Cell( + why="the 48-token maximum with the NVFP4 activation", + spec=dict(m=48, k=6144, fp4=True), + ), + # Every other row count the router takes (modeling's ROUTER_TC), with and without the NVFP4 activation. + *( + Cell( + why=f"{m} rows: a router row count between the ones above", + spec=dict(m=m, k=6144, fp4=m % 2 == 1), + ) + for m in range(2, 49) + if m not in (2, 8, 13, 48) + ), + ) + + note = """ + The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both + sum the same 16 slice partials in the same order. Only the one the test device picks is driven. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> tuple[torch.Tensor, ...]: + """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" + with true_fp32_matmul(): + logits = x.float() @ tf32_rna(w).t() + if sf_scale is None: + return (logits,) + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return logits, packed, sf + + def is_valid( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + if sf_scale is not None: + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/__init__.py new file mode 100644 index 000000000000..acec8fdbb26f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalization entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..9093d53b3ce7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step embedding gather and the first layer's input RMSNorm in one CuTe DSL launch, into caller buffers.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm + +__all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] + +VOCAB = 200064 # MiniMax-M3's vocabulary; the decode table is the full (replicated) embedding + + +class M3EmbeddingFirstnorm(OpWrapper): + """Write ``table[ids]`` into ``raw`` (bf16 [M, 6144], a zero row for an id outside ``[0, V)``) and + ``raw * rsqrt(mean(raw^2) + eps) * weight`` into ``out``; returns ``(out, raw)`` -- the first layer's normed input + and its residual. ``ids`` int32 or int64 [M], M = 1..16; ``table`` bf16 [V, 6144] contiguous; ``weight`` bf16 + [6144]. The first call per (M, id dtype, table shape, PDL) compiles (``DecodeEmbedding`` does it at build).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target builds every (rows 1..16, id dtype) owner; the kernel is one CTA per row, so the rows differ only in + # grid size: the ends of the range and both id dtypes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, int32 ids", spec=dict(rows=1, id_dtype=torch.int32)), + Cell(why="eight requests, int32 ids", spec=dict(rows=8, id_dtype=torch.int32)), + Cell(why="one request, int64 ids", spec=dict(rows=1, id_dtype=torch.int64)), + Cell( + why="six rows (one request verifying a draft), int64 ids", + spec=dict(rows=6, id_dtype=torch.int64), + ), + # Every other row count the decode dispatch takes (modeling's EMBED_FIRSTNORM), both id dtypes. + *( + Cell( + why=f"{rows} rows: a decode row count between the ones above", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in (2, 3, 4, 5, 7) + ), + # B5, the batch-4 decode steps (modeling's EMBED_FIRSTNORM past 8 rows): every row count 9..16, both id dtypes. + *( + Cell( + why=f"{rows} rows (batch-4 decode step)", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in range(9, 17) + ), + ) + + note = """ + The gathered rows are exact. The norm is flashinfer's RMSNorm arithmetic in its reduction order with a fast + rsqrt, where `reference` sums in fp64 and takes torch's rsqrt: within the band the test derives, not bit for bit. + SM107 against the stock embedding + RMSNorm module (flashinfer's kernel, 768 threads): the gathered rows are + bitwise equal; the normed rows differ in one bf16 ulp on ~5e-6 of the elements (~1.6% of rows, 1-2 elements each: + the fp32 sum-of-squares tree differs, so the rsqrt input can differ in its last bits). + An id outside [0, V) gives a zero row and a zero norm (torch's embedding would raise). No accepted input is known + to compute wrongly: the host side rejects any other dtype, shape or a strided buffer. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return m3_embed_norm(ids, table, weight, eps, raw, out, use_pdl) + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x * rstd * weight.double()).to(torch.bfloat16), rows + + +m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/__init__.py new file mode 100644 index 000000000000..d25c41dc9557 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries this target alone calls: the Eagle3 verify ops, the capture / fc_norm RMSNorms and +the step-input staging.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..a1993924a112 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_chunk_norm import rmsnorm_rows + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{4 * n} verify rows of {n} requests (b4 / b8), layer slot {n % 3}, residual added", + spec=dict(m=4 * n, slot=n % 3, residual=True), + ) + for n in range(3, 8) + ), + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ) + + note = """ + fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional + one-ulp move against torch's statistic. + """ + + def __call__( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) + return dst + + def reference( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" + x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h + return rmsnorm_rows(x, w, eps) + + def is_valid( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> None: + # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read + # past its elements without an error. + assert w.is_contiguous(), "w must be dense [H]" + + +m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..453edec27764 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states +(``m3_spec_tail.chunk_norm``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{n} requests' {4 * n} verify rows (b4 / b8), dense [{4 * n}, 3 x 6144]", + spec=dict(m=4 * n, c=3, h=6144, extra_cols=0), + ) + for n in range(3, 8) + ), + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="two requests' 8 verify rows (b2) from the capture buffer, dense [8, 3 x 6144]", + spec=dict(m=8, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ) + + note = """ + fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then + against torch's statistic; the cells' band states it. + """ + + def __call__( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) + return dst + + def reference( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" + c, h = w.shape + x = src[:, : c * h].reshape(src.shape[0], c, h) + return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) + + def is_valid( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ) -> None: + c, h = w.shape + assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + + +m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..4fe14d538a5f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat +(``m3_spec_tail.draft_input``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3DraftInput(OpWrapper): + """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = + hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of + a vocab-local embedding). Returns ``(out, residual)``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling's DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the + # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests + # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", + spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", + spec=dict(m=n, h=6144, gather=True, invalid_id=False), + ) + for n in range(1, 9) + ), + Cell( + why="third pass of 3 requests: one row each, hidden read in place", + spec=dict(m=3, h=6144, gather=False, invalid_id=False), + ), + Cell( + why="one row with an out-of-table id: the embedding half is zero", + spec=dict(m=1, h=6144, gather=False, invalid_id=True), + ), + ) + + note = """ + The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's + summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, + also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the + drafter's embedding is the whole table, so its ids are always in range. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl + ) + return out, residual + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / + ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" + valid = (ids >= 0) & (ids < table.shape[0]) + emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) + emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) + hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] + ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) + return ref, (hid.clone() if residual is not None else None) + + def is_valid( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> None: + h = table.shape[1] + # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. + assert table.is_contiguous() and hidden.is_contiguous(), ( + "table / hidden must be dense [*, H]" + ) + if rows is not None: + assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" + else: + assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" + # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + assert residual is None or ( + residual.is_contiguous() and residual.shape == (ids.shape[0], h) + ), "residual must be contiguous [M, H]" + + +m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py new file mode 100644 index 000000000000..e322e1787230 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter MLP's plain SwiGLU as one C++ launch (modeling's ``DRAFT_SILU_CPP``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3DraftSilu(OpWrapper): + """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: + trtllm::minimax_m3_sm107_tp2_silu_and_mul (minimaxM3Sm107Tp2/siluAndMul.cu), bit for bit the Triton + ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP + runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs + the Triton op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP2 drafter MLP shard: intermediate 18432 / 2 = 9216, gate | up 18432 wide, at the drafter's decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP2 drafter MLP at 1 row", spec=dict(rows=1, d=9216)), + Cell(why="the TP2 drafter MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=9216)), + Cell(why="the TP2 drafter MLP at 8 rows", spec=dict(rows=8, d=9216)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (the same PTX instruction sequence), checked over every + bf16 gate and every bf16 up pattern. + """ + + def __call__(self, x: torch.Tensor) -> Optional[torch.Tensor]: + if ( + x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 16 != 0 + or x.data_ptr() % 16 != 0 + or x.shape[0] > 65535 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp2_silu_and_mul(x, out) + return out + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul(x) + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_draft_silu = M3DraftSilu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..708e6c5caa71 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py @@ -0,0 +1,432 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One pass of the fused greedy Eagle3 draft loop: the drafter's vocab-sharded BF16 head (streamed from its lossless +12-bit coding, ``ops.gemm.m3_target_head``, or as BF16 fragments), the cross-rank argmax and the loop's per-pass +bookkeeping in one kernel (``m3_draft_step.draft_step``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ..gemm.m3_target_head import decode as decode_coded +from .m3_shard_accept import exchange_tokens, local_argmax, make_keys + +#: Requests (head rows) per pass, and exchange rows per call half of the workspace. +MAX_TOKENS = 8 +#: 16-row weight tiles one CTA can own (the kernel's shared-memory accumulator). +MAX_TILES_PER_CTA = 32 +#: Draft length of the fused loop, and the drafter's TP2 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 100032 +HIDDEN = 6144 + + +def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the + 12-bit coded stream with its exception lists decoded.""" + if exc_off is None: + return from_fragments(wf) + n = (exc_off.numel() - 1) * 16 + return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + + +def shard_logits( + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() + + +class M3DraftStep(OpWrapper): + """One fused greedy draft pass over ``rows <= 8`` requests (one head row each). + + Head input ``x[x_rows[b]]`` (pass 0: the last accepted verify row of each request) or ``x[b]``. The head is this + rank's BF16 lm_head shard in fragment order (``m3_draft_step.bf16_fragments``, global ids ``vocab_offset ..``), + or with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base`` its lossless 12-bit coding (``wf`` = the uint8 stream + of ``m3_target_head.encode_chunked``, which must round-trip to the BF16 shard); each logit is the fp32 dot product + rounded to BF16 and the row's token is the global argmax over all ranks' + shards (ties, -0 / +0 included, to the lowest global id), exchanged through the multicast slots + ``[2][8][world]``. The last CTA then writes ``draft_out[b, step]`` and ``ids_next[b]`` (the token), ``pos_next[b] + = pos_src[x_rows[b] or b] + 1``, and advances the draft KV lengths as the in-tree draft loop does between passes + (pass 0 with ``num_acc``: ``kv_lens -= draft_len - num_acc`` and ``seq_lens = 1``; later passes: ``kv_lens += + 1``). ``out_keys`` (optional) receives each row's winning key. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling's DRAFT_STEP): draft_len 3 at the TP2 GEN worker's 1..2 requests, + # the drafter's full-vocabulary BF16 head (100032 rows per rank of 2, hidden 6144; 6252 tiles, 30 per CTA on 212 + # SMs). Pass 0 reads the verify rows' hidden (4 per request) through x_rows and rolls the KV back; the later passes + # read one row per request. BF16 fragments only: on sm_107 the 12-bit coded head is slower than the BF16 stream + # (modeling's DRAFT_HEAD_CODED is False). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"pass 0 at {n} request(s): head rows gathered from {4 * n} verify rows, KV rollback, seq_lens = 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=0, + rank=n % 2, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"pass 1 at {n} requests (b4 / b8): one row each, kv_lens + 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=1, + rank=n % 2, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ) + for n in (4, 8) + ), + Cell( + why="pass 2 at 8 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=8, + draft_len=DRAFT_LEN, + step=2, + rank=0, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + coded=False, + ), + ), + Cell( + why="pass 1 at two requests: one row each, kv_lens + 1, rank 1's ids", + spec=dict( + nreq=2, + draft_len=DRAFT_LEN, + step=1, + rank=1, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ), + Cell( + why="last pass (step 2) at one request, rank 0", + spec=dict( + nreq=1, + draft_len=DRAFT_LEN, + step=2, + rank=0, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ), + Cell( + why="pass 2 at 2 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=2, + draft_len=DRAFT_LEN, + step=2, + rank=1, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + coded=False, + ), + ), + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit; + the coded head's weight is the BF16 shard bit for bit; its exceptions decode to +0 in the MMAs and their exact + products are added at the tile's flush) + and the argmax is the in-tree one over BF16 logits, so a token differs from the in-tree path's only where the fp32 + sum, accumulated in another order, rounds to the other side of a BF16 boundary and that decides the row's maximum. + The cells' rows are decidable and the peers' keys sit 0.25 away (or tie exactly); the edge tests (forced BF16 + ties, -0 / +0, top-2 one fp32 step apart that round to one BF16 value, adjacent BF16 values) use integer data + whose fp32 sums are exact in any order and hold the kernel to the in-tree path bit for bit. The cross-rank + exchange is certified on one device with the peers' keys written into the slots (see m3_shard_accept). + """ + + def __call__( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + bnd: Optional[torch.Tensor] = None, + bnd_cnt: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_sm107tp2_draft_step( + x, + x_rows, + rows, + wf, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + num_ctas, + pdl, + out_keys, + exc_off, + exc_pos, + exc_w, + exp_base, + tile_lo, + tile_hi, + total_ctas, + bnd, + bnd_cnt, + ) + + def reference( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + bnd: Optional[torch.Tensor] = None, + bnd_cnt: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Every buffer the pass writes, computed on copies: "logits" (this rank's BF16 logits as fp32), "local_keys" + (this rank's key per row), "keys" (the winning key), "tokens", "draft_out", "ids_next", "pos_next", "kv_lens" + and "seq_lens". + + Reads ``ws_uc``, ``counters`` and ``kv_lens`` before the call. ``keys``, ``ws_mc_ptr``, ``num_ctas``, ``pdl`` + and ``out_keys`` do not change the result. + """ + h = x[x_rows.long()] if x_rows is not None else x[:rows] + logits = shard_logits(h, wf, exc_off, exc_pos, exc_w, exp_base) + best, idx = local_argmax(logits, vocab_offset) + local_keys = make_keys(best, idx) + tokens = exchange_tokens(local_keys, ws_uc, counters, rank, world, MAX_TOKENS) + half = int(counters[1].item()) & 1 + slots = ws_uc.view(torch.int64)[ + half * MAX_TOKENS * world : (half * MAX_TOKENS + rows) * world + ] + slots = slots.view(rows, world).clone() + slots[:, rank] = local_keys + sign = torch.iinfo(torch.int64).min + win = (slots ^ sign).max(dim=1).values ^ sign + d_out = draft_out.clone() + d_out[:rows, step] = tokens + src = x_rows.long() if x_rows is not None else torch.arange(rows, device=x.device) + lens = kv_lens.clone() + sl = seq_lens.clone() if seq_lens is not None else None + if step == 0 and num_acc is not None: + lens[:rows] -= draft_len - num_acc[:rows] + sl[:rows] = 1 + elif step > 0: + lens[:rows] += 1 + out = { + "logits": logits, + "local_keys": local_keys, + "keys": win, + "tokens": tokens, + "draft_out": d_out, + "ids_next": tokens.clone(), + "pos_next": pos_src[src] + 1, + "kv_lens": lens, + } + if sl is not None: + out["seq_lens"] = sl + return out + + def is_valid( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + bnd: Optional[torch.Tensor] = None, + bnd_cnt: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[:rows].any()), "keys must start zeroed" + # Equal-chunk split: the shared-tile partials and arrival counts start (and are left) zero. + assert (bnd is None) == (bnd_cnt is None), "bnd and bnd_cnt come together" + if bnd is not None: + assert not bool(bnd.any()) and not bool(bnd_cnt.any()), ( + "bnd / bnd_cnt must start zeroed" + ) + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + # The kernel reads the head row x + x_rows[b] * K and the position pos_src[x_rows[b]] with no bound. + if x_rows is not None: + assert bool(((x_rows[:rows] >= 0) & (x_rows[:rows] < x.shape[0])).all()), ( + "x_rows must index x" + ) + assert bool((x_rows[:rows] < pos_src.shape[0]).all()), ( + "pos_src must cover the head rows" + ) + else: + assert pos_src.shape[0] >= rows, "pos_src must cover the head rows" + # Each CTA accumulates at most MAX_TILES_PER_CTA tiles in shared memory. + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + assert -(-tiles // num_ctas) <= MAX_TILES_PER_CTA, "too many weight tiles per CTA" + if exc_off is not None: + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), ( + "exc_off must end at the exception count" + ) + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + # A NaN logit maps to a key above every finite one and wins the row silently. + h = x[x_rows.long()] if x_rows is not None else x[:rows] + assert not bool(torch.isnan(h).any()), "NaN head inputs win the argmax" + if num_acc is not None: + na = num_acc[:rows] + assert bool(((na >= 1) & (na <= draft_len + 1)).all()), ( + "num_acc counts must lie in [1, draft_len + 1]" + ) + + +m3_draft_step = M3DraftStep() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..2001bf183b54 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.m3_force_vec import force_vec_kernel, force_vec_req_kernel + + +class M3ForceVec(OpWrapper): + """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` + (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 + ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + force_vec_kernel[(1,)]( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + """The counts (a new tensor; ``counter`` and ``out`` are not written).""" + c = int(counter.item()) + 1 + slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) + idx = (c * counter_stride + slots * slot_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec = M3ForceVec() + + +class M3ForceVecReq(OpWrapper): + """For each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(keys[b] * key_stride) & pool_mask] < frac), max_total)`` into ``out`` (int32): + one draw per request at its own key (its KV length at the verify), no shared counter. ``block`` is the launch + width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="two requests, fraction 0.3", spec=dict(batch=2, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + force_vec_req_kernel[(1,)]( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + """The counts (a new tensor; ``out`` is not written).""" + k = keys[:batch_size].to(torch.int64) + r = pool[(k * key_stride) & pool_mask] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec_req = M3ForceVecReq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..34fb3387b996 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampled (temperature + top-p) Eagle3 verification over the TP-sharded vocabulary in one kernel +(``m3_sampled_accept.sampled_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_shard_accept import strict_accept + +#: Verify rows per call. +MAX_ROWS = 48 +#: The fused sampling kernel's 1 / T guard. +TEMP_EPSILON = 1e-6 + + +def nucleus_probs( + logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +class M3SampledAccept(OpWrapper): + """Sampled verification of ``rows = nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard. + + Each row draws one token from its top-p nucleus of softmax(logits / T) over the full (all-rank) vocabulary by + distributed rejection sampling (Philox keyed by the row's seed / offset, the row, the attempt and the CTA), then + the strict acceptance runs against ``draft`` as in ``m3_shard_accept``: ``acc_tokens``, ``num_acc``, + ``gather_ids``, optional ``next_new[:, 0]``. ``tokens`` receives the sampled token per row; ``extra_attempts`` + accumulates exchange rounds beyond the second. The exchange goes through the multicast workspace (two call + halves of ``minimax_m3_sampled_accept_ws_words(world)`` words) with the (finished rows, call) ``counters``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3 (T = 4 rows per request) at the TP2 GEN worker's batch: 1 / 2 requests (4 / 8 rows, 8 CTAs per + # row), and the 48-row maximum (4 CTAs per row, 50 KB of staged logits per CTA); the 100032-column target shard + # (vocab 200064 / 2). Single rank (see note). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"b4 / b8 ({4 * n} rows, {n} requests): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=n, + vs=100032, + vocab_offset=100032 * (n % 2), + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ) + for n in range(3, 9) + ), + Cell( + why="b8 (32 rows, 8 CTAs per row = 256 CTAs): AgentX sampling, the draws follow the nucleus", + spec=dict( + nreq=8, + vs=100032, + vocab_offset=0, + check="distribution", + top_p=0.95, + temperature=1.0, + force="total", + ), + ), + Cell( + why="b1 (4 rows), top_p 0.8 over a 12-token head: the draws follow the nucleus distribution", + spec=dict( + nreq=1, vs=100032, vocab_offset=0, check="distribution", top_p=0.8, temperature=0.7 + ), + ), + Cell( + why="b2 (8 rows): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=2, + vs=100032, + vocab_offset=100032, + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ), + Cell( + why="b2 (8 rows) with top_p = 1 (no cut: the first draw stands), next_new and forced counts", + spec=dict( + nreq=2, + vs=100032, + vocab_offset=100032, + check="distribution", + top_p=1.0, + temperature=1.0, + force="vec", + ), + ), + Cell( + why="AgentX sampling (T 1.0, top_p 0.95) at b1 with forced counts: the draws follow the nucleus", + spec=dict( + nreq=1, + vs=100032, + vocab_offset=0, + check="distribution", + top_p=0.95, + temperature=1.0, + force="total", + ), + ), + Cell( + why="48 rows (4 CTAs per row): the second launch geometry, nucleus distribution at top_p 0.5", + spec=dict( + nreq=12, + vs=100032, + vocab_offset=0, + check="distribution", + top_p=0.5, + temperature=1.3, + force="total", + ), + ), + ) + + note = """ + What is sampled is gated statistically, the rest exactly. Every drawn token must lie in the reference nucleus; + the empirical distribution over many calls (distinct Philox offsets) must sit within a stated total-variation + band of the reference's; the acceptance outputs must equal the strict acceptance of the drawn tokens exactly. + The Philox stream itself is not reproduced in torch: the kernel's draw uses the hardware exp / log + approximations (__expf, __logf) and a CTA-order inverse CDF, which torch does not reproduce bit for bit. + + Certified at world 1 only. The decision is replicated from every rank's per-round records, which another rank + computes from its own shard on the fly, so peers cannot be stood in for on one device; the cross-rank path + needs the multi-device job and is not covered here. Rows with top-k or min-p must not reach this op (it applies + neither); is_valid states it. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept( + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + rows = logits.shape[0] + # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the + # wrong rows. + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + # Per-row parameters are read at the row index, unchecked. + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), ( + "per-row parameters must cover every row" + ) + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + # The finished-row counter elects the CTA that runs the acceptance. + assert int(counters[0].item()) == 0, ( + "counters[0] (finished rows) must be zero between calls" + ) + + +m3_sampled_accept = M3SampledAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..13ec40116191 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept_ws_words.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Size of the sampled-verify exchange workspace (``m3_sampled_accept.SampledAcceptWorkspace.create``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Exchange rounds per row, verify rows per call, CTAs per row (at most), uint64 words per record. +MAX_ATTEMPTS = 16 +MAX_ROWS = 48 +MAX_SPLITS = 8 +RECORD_WORDS = 4 + + +class M3SampledAcceptWsWords(OpWrapper): + """uint64 words of the multicast workspace ``minimax_m3_sm107tp2_sampled_accept`` exchanges through for + ``world`` ranks: two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="the tp2 target's workspace", spec=dict(world=2)), + Cell(why="a single rank (local workspace, no multicast)", spec=dict(world=1)), + ) + + note = """ + A host computation (no device work); it is an op so the workspace size and the kernel's indexing come from one + definition. + """ + + def __call__(self, world: int) -> int: + return torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept_ws_words(world) + + def reference(self, world: int) -> int: + """Two halves of attempts x rows x (world x splits) records x record words.""" + return 2 * MAX_ATTEMPTS * MAX_ROWS * world * MAX_SPLITS * RECORD_WORDS + + +m3_sampled_accept_ws_words = M3SampledAcceptWsWords() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..408759009472 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Greedy Eagle3 verification from this rank's logits shard: per-row argmax, the cross-rank (value, index) exchange +and the strict acceptance in one kernel (``m3_shard_accept.shard_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Verify rows per call, and exchange rows per call half of its workspace. +SHARD_ROWS = 48 + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out + + +class M3ShardAccept(OpWrapper): + """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` + (global ids ``vocab_offset ..``). + + Each row's token is its global argmax over all ranks' shards: keys ``ordered(value) << 32 | ~index`` (ties take + the lowest index, as ``torch.argmax``) are max-reduced over the shard, pushed into every rank's exchange slot + through the multicast slots ``[2][48][world]`` (Lamport: 0 = empty) and max-reduced over the ranks by the last + CTA, which then runs the strict acceptance: ``acc_tokens [nreq, D + 1]``, ``num_acc [nreq]`` (``force_vec`` / + ``force_total`` override the counts), ``gather_ids [nreq]`` (the first draft pass's row per request) and, when + given, ``next_new[:, 0]`` (the last accepted token). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every greedy verify step of the target (modeling's SHARD_ACCEPT): draft_len 3 at the TP2 GEN worker's 1..2 + # requests (4..8 rows), plus the 48-row maximum; the target lm_head's BF16 shard logits (100032 columns per rank + # at TP2). The logits come from the head GEMM, so the op only reduces. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"b4 / b8: {4 * n} verify rows ({n} requests), rank {n % 2} of 2", + spec=dict( + nreq=n, draft_len=3, vs=100032, world=2, rank=n % 2, force="none", strided=False + ), + ) + for n in range(3, 9) + ), + Cell( + why="b1: 4 verify rows (1 request x T = 4), rank 0 of 2, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=100032, world=2, rank=0, force="none", strided=False), + ), + Cell( + why="b2: 8 verify rows, rank 1, global ids offset by 100032", + spec=dict(nreq=2, draft_len=3, vs=100032, world=2, rank=1, force="none", strided=False), + ), + Cell( + why="8 rows (2 requests) read from a row-strided view of a wider logits buffer, forced counts", + spec=dict(nreq=2, draft_len=3, vs=100032, world=2, rank=0, force="vec", strided=True), + ), + # W2 (WIDE_SPEC_TAIL): 13..48 requests (52..192 verify rows) as whole-request chunks of 12 (the 48-row + # geometry), each chunk a launch on its own rows; world 1 (the chunked exchange is the per-launch one above). + *( + Cell( + why=f"wide: {4 * n} verify rows ({n} requests) in {-(-n // 12)} chunks, world 1", + spec=dict( + nreq=n, + draft_len=3, + vs=100032, + world=1, + rank=0, + force="vec" if n % 3 == 0 else "none", + strided=n % 5 == 0, + ties=n % 2 == 0, + ), + ) + for n in range(13, 49) + ), + # W2b: the drafter's greedy top-1 (draft length 0: one row per request, token = the global argmax), 1..48 + # requests in one launch and 49..56 in two (spec_worker.greedy_sample_draft_with_tp_gather). + *( + Cell( + why=f"draft top-1: {n} rows (draft length 0), world 1", + spec=dict( + nreq=n, draft_len=0, vs=100032, world=1, rank=0, force="none", strided=n == 20 + ), + ) + for n in (1, 8, 20, 46, 48, 56) + ), + # 9..11 requests (36..44 rows): one launch each (the 12-request cell below is the 48-row one). + *( + Cell( + why=f"wide: {4 * n} verify rows ({n} requests), one launch, world 1", + spec=dict( + nreq=n, draft_len=3, vs=100032, world=1, rank=0, force="none", strided=False + ), + ) + for n in range(9, 12) + ), + Cell( + why="the 48-row maximum (12 requests at draft 3), force_total, rows with exact cross-rank ties", + spec=dict( + nreq=12, + draft_len=3, + vs=100032, + world=2, + rank=1, + force="total", + strided=False, + ties=True, + ), + ), + ) + + note = """ + The op reduces BF16 logits exactly, so the reference is bit-exact, including cross-rank ties (lowest global + index wins). The cross-rank exchange is certified on one device: the other ranks' keys are written into this + rank's slots before the call, which is what their multicast pushes would have put there; the multicast mapping + itself (ws_mc_ptr = this rank's slots, as at world 1) is not exercised across devices here. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + rows = logits.shape[0] + if rows > SHARD_ROWS: + # W2: whole requests per launch; batch_idx is the identity layout (is_valid), so each chunk passes its + # local rows and its gather ids are shifted back to the full verify layout. + k1 = draft.shape[1] + 1 + per = SHARD_ROWS // k1 + nreq = draft.shape[0] + local = torch.arange(per, dtype=torch.int32, device=logits.device) + for c0 in range(0, nreq, per): + c1 = min(nreq, c0 + per) + self( + logits[c0 * k1 : c1 * k1], vocab_offset, keys, counters, ws_uc, ws_mc_ptr, rank, world, + draft[c0:c1].contiguous(), local[: c1 - c0], acc_tokens[c0:c1], num_acc[c0:c1], + gather_ids[c0:c1], force_total, None if next_new is None else next_new[c0:c1], num_ctas, + None if force_vec is None else force_vec[c0:c1], + ) # fmt: skip + gather_ids[c0:c1] += c0 * k1 + return + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + num_ctas, + force_vec, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Lowest-index maximum of each row's BF16 logits, the largest key among it and the peers' keys in + ``ws_uc``, then the strict acceptance. Reads ``ws_uc`` / ``counters`` before the call; the scratch and + output buffers and ``num_ctas`` do not change the result.""" + best, idx = local_argmax(logits.float(), vocab_offset) + if world == 1: + tokens = idx.to( + torch.int32 + ) # no peers: the exchange returns this rank's own pick (any row count) + else: + tokens = exchange_tokens(make_keys(best, idx), ws_uc, counters, rank, world, SHARD_ROWS) + out = strict_accept(tokens, draft, batch_idx, force_total, next_new, force_vec) + out["tokens"] = tokens + return out + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[: logits.shape[0]].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + assert ws_uc.numel() >= 2 * SHARD_ROWS * world, "ws_uc must hold [2][48][world] slots" + # A NaN logit maps to a key above every finite one and wins the row silently. + assert not bool(torch.isnan(logits).any()), "NaN logits win the argmax" + nreq = draft.shape[0] + assert bool(((batch_idx[:nreq] >= 0) & (batch_idx[:nreq] < nreq)).all()), ( + "batch_idx must index the nreq requests" + ) + if logits.shape[0] > SHARD_ROWS: + # The chunked (W2) call relocates each chunk's rows: only the identity layout maps back. + assert torch.equal(batch_idx[:nreq].cpu(), torch.arange(nreq, dtype=batch_idx.dtype)), ( + "a chunked verify needs the identity batch_idx layout" + ) + if force_vec is not None: + fv = force_vec[:nreq] + assert bool(((fv >= 1) & (fv <= draft.shape[1] + 1)).all()), ( + "force_vec counts must lie in [1, D + 1]" + ) + + +m3_shard_accept = M3ShardAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..ff07c54de1f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> +``m3_spec_tail.slot_scatter``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: + """src rows zero-padded (or cut) to ``width`` columns.""" + out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) + w = min(width, src.shape[1]) + out[:, :w] = src[:, :w] + return out + + +class M3SlotScatter(OpWrapper): + """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, + ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The + destinations are the store's transposed [slots, width] views.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of + # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) + # is zero-padded. + CELLS: tuple[Cell, ...] = ( + Cell( + why="8 rows into shuffled slots, widths equal to the store's", + spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), + ), + Cell( + why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", + spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), + ), + ) + + note = """ + Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. + """ + + def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): + torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) + + def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): + """The four destinations after the call, on copies.""" + s = slots.long() + out = {} + for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): + r = dst.clone() + r[s] = padded_rows(src[: s.numel()], dst.shape[1]) + out[name] = r + lr = dl.clone() + lr[s] = sl[: s.numel()] + out["dl"] = lr + return out + + def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): + assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( + "slots must index the store" + ) + # Two rows into one slot race (which one lands is unordered). + assert slots.unique().numel() == slots.numel(), "slots must be distinct" + assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + + +m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py new file mode 100644 index 000000000000..8925fc8ab5b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's speculative step correction as one launch (``RuntimeHooks.spec_offsets_op``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger + +# The engagement line is logged once per process (module state: __call__ reads no instance state). +_ENGAGED: list = [] + + +class M3SpecOffsets(OpWrapper): + """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp2_spec_offsets (minimaxM3Sm107Tp2/specOffsets.cu), the + same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other + layouts (the caller runs the ATen adds).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="1 request, 4 tokens (draft length 3)", spec=dict(tokens=4, seqs=1)), + Cell(why="8 requests, 32 tokens", spec=dict(tokens=32, seqs=8)), + Cell(why="wrap-around values", spec=dict(tokens=8, seqs=2, wrap=True)), + ) + + note = "Integer adds: the reference (the two in-place ATen adds) is exact." + + def __call__(self, pos, pos_off, kv, kv_off) -> bool: + ts = (pos, pos_off, kv, kv_off) + if ( + any( + t.dtype != torch.int32 or t.dim() != 1 or t.stride(0) != 1 or not t.is_cuda + for t in ts + ) + or pos.numel() != pos_off.numel() + or kv.numel() != kv_off.numel() + ): + return False + torch.ops.trtllm.minimax_m3_sm107_tp2_spec_offsets(pos, pos_off, kv, kv_off) + if not _ENGAGED: + _ENGAGED.append(True) + logger.info("[M3 host] spec step correction: one launch taken (spec_offsets_op)") + return True + + def reference(self, pos, pos_off, kv, kv_off) -> None: + pos += pos_off + kv += kv_off + + def is_valid(self, pos, pos_off, kv, kv_off) -> None: + assert pos.numel() == pos_off.numel() and kv.numel() == kv_off.numel(), ( + "equal lengths per pair" + ) + + +m3_spec_offsets = M3SpecOffsets() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..a0da8b337574 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch +(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS: tuple[Cell, ...] = ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + Cell( + why="an unchanged batch of 3 previous requests (c3) in shuffled slots", + spec=dict(p=3, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="an unchanged batch of 4 previous requests (c4) in shuffled slots", + spec=dict(p=4, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + ) + + note = """ + Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid + states them. + """ + + def __call__( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ) + + def reference( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" + nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) + s = slots[:P].long() + ids_r, draft_r = ids.clone(), draft.clone() + ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) + draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) + pos_r = torch.zeros_like(pos_off) + pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] + kv_r = torch.zeros_like(kv_off) + kv_r[b0 : b0 + P] = lens[s] - T + return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} + + def is_valid( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or + # writes another buffer. + s = slots[:P] + assert bool( + ( + (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) + ).all() + ), "slots must index the store" + assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( + "ppi must index lens" + ) + assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" + assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" + assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( + "offset buffers too short" + ) + + +m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py new file mode 100644 index 000000000000..080c3018e4b6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode graph's step-input prologue in one launch (this target's ``RuntimeHooks.pack_prologue_op``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepPrologue(OpWrapper): + """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch + (trtllm::minimax_m3_sm107_tp2_step_prologue, minimaxM3Sm107Tp2/stepPrologue.cu). + + ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one + ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the + window, a device copy's a device address). The launch writes up to ``max_desc`` copies of each kind, every host + copy before any device copy, then zeroes the header; with ``epoch`` it increments ``epoch[0]`` once. The default + prologue does the same with two parallel-scatter launches (``step_inputs._graph_scatter_kernel``) and a header + fill.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP2 decode step's packs: a few dozen host copies of 4 to a few hundred bytes (positions, KV lengths, slots, + # the spec sampler's uploads), a handful of device copies (the graph runner's static inputs), some reading what a + # host copy of the same window wrote. + CELLS: tuple[Cell, ...] = ( + Cell(why="an empty window (a replay with no new upload)", spec=dict(h2d=0, d2d=0)), + Cell(why="host copies only, odd sizes and alignments", spec=dict(h2d=40, d2d=0)), + Cell( + why="host copies feeding device copies (a device copy reads a host copy's destination)", + spec=dict(h2d=24, d2d=8, chained=True), + ), + Cell( + why="every descriptor slot (128 + 128 of the pack's 256)", spec=dict(h2d=128, d2d=128) + ), + ) + + note = """ + Byte copies: the reference (the default prologue's Triton launches and fill on the same window) is exact. + Destinations are raw addresses; is_valid cannot check them. + """ + + def __call__(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + torch.ops.trtllm.minimax_m3_sm107_tp2_step_prologue(window, epoch, max_desc) + + def reference(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + """The default prologue on the same arguments: it writes the same destinations.""" + from tensorrt_llm._torch.pyexecutor.step_inputs import _graph_scatter_kernel + + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + bump = epoch is not None and not device_src + _graph_scatter_kernel[(max_desc,)]( + window, + stage_i64, + epoch if bump else stage_i64, + DEVICE_SRC=device_src, + BLOCK=256, + BUMP_EPOCH=bump, + num_warps=1, + ) + window[:16].zero_() + + def is_valid(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + assert window.dtype == torch.uint8 and window.is_contiguous(), "window: contiguous uint8" + assert window.numel() >= 16 + 24 * max_desc, "window: header and max_desc descriptors" + counts = window[:16].view(torch.int64).tolist() + assert counts[0] + counts[1] <= max_desc, "at most max_desc copies in all" + + +m3_step_prologue = M3StepPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/weights.py new file mode 100644 index 000000000000..fa0bc1ba7208 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/weights.py @@ -0,0 +1,865 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_107 / tp2. + +Two tables, one per checkpoint the target reads: + +* ``MODEL_WEIGHTS`` (``MiniMaxM3Tp2Weights``): the vision-language checkpoint's language model, + ``language_model.*``. Every rank is handed the whole checkpoint, so every tensor-parallel split is this + table's: the Q heads, the K/V/index-Q heads (two each per rank), the dense MLPs and the shared expert's + intermediate, the o_proj / down-projection input columns, every routed expert's intermediate (all 128 + experts live on every rank), and the vocabulary of the shell's ``lm_head``. The index-K head is replicated. +* ``DRAFT_WEIGHTS`` (``MiniMaxM3Tp2DraftWeights``): the one-layer BF16 Eagle3 drafter's own checkpoint + (``speculative_model``), loaded through the shell's ``load_draft_weights`` when speculation is on. + +Storage is HF ``[out, in]`` row-major. Five families need a transform, each a layout a kernel demands: + +* **MXFP8 block scales** (every attention projection, the dense MLPs and the shared expert): the checkpoint + stores UE8M0 bytes ``[N, K / 32]``; the GEMMs read the 128x4-swizzled order + (``trtllm::block_scale_interleave``). A fused projection concatenates its shards' scales first and swizzles + once. +* **The five-way sparse-layer projection** ``[Q | K | V | index-Q | index-K]`` and the three-way dense one + ``[Q | K | V]``: the per-rank shards concatenated along N. +* **Gated MLPs** (dense and shared expert): ``[gate ; up]`` rows. +* **NVFP4 routed experts** (layers 3..59): ``[up (w3) ; gate (w1)]`` rows, then the trtllm-gen row shuffle of + the weight and of its E4M3 scales (gated-activation reorder + 128-row epilogue tile shuffle, + ``trtllmgen_maybe_get_cached_*_permute_indices``) and the 128x4 swizzle of the scales. +* **Gemma boundary norms**: ``input_layernorm``, ``post_attention_layernorm`` and ``model.norm`` are stored as + Gemma norms, ``(1 + w) * x``; the fused all-reduce + RMSNorm kernels apply ``w * x``, so ``1`` is added + (in BF16) at load. The per-head Q/K/index norms stay Gemma: their kernels apply the offset themselves. + +The routed experts' per-tensor scalars (input scales and second-level weight scales, every expert) load raw; +``modeling`` derives the kernels' global scales, alphas and activation constants from them once the weights +are real. + +Loading contract: ``load(model, weights)`` consumes the engine-provided checkpoint dict, fills every declared +parameter exactly once and asserts coverage in both directions: every declared parameter written, and every +checkpoint key either consumed or in ``expected_unconsumed`` (the vision tower and its projector, which this +text-only target does not build, and the language model's multi-token-prediction head). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) + +#: ``layers=ALL_LAYERS`` is the common case: a weight every decoder layer carries. +ALL_LAYERS = object() + +_LM = "language_model." +#: Checkpoint key families this target does not build: the vision tower and its two projectors (text only), and the +#: language model's multi-token-prediction head (speculation, when on, is the separate Eagle3 drafter). +_UNBUILT_PREFIXES = ( + "vision_tower.", + "multi_modal_projector.", + "patch_merge_mlp.", + _LM + "model.mtp.", +) + +MX_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 elements along K +NV_BLOCK = 16 # NVFP4: one E4M3 scale per 16 elements along K +EPILOGUE_TILE_M = ( + 128 # the trtllm-gen MoE runner's row tile: fixes its weight and scale row shuffles +) + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + ``shape`` takes the dimension bundle and returns the declared shape. ``src`` is a callable + ``(d, i) -> [(key or key tuple, index into the param | None), ...]`` with keys relative to the checkpoint + root; a key tuple hands the transform those tensors together. ``layers`` is ``ALL_LAYERS``, ``None`` (not + per-layer) or ``d -> layer indices``. ``transform(src, d)`` runs on the destination's device before the copy. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + +# ------------------------------------------------------------------------------------------------ shards + + +def _col( + t: torch.Tensor, d: SimpleNamespace, tp: int | None = None, rank: int | None = None +) -> torch.Tensor: + """This rank's rows (column-parallel output split), even split, sliced on the host and moved to the + destination device (only this rank's bytes are read).""" + tp = d.tp if tp is None else tp + rank = d.rank if rank is None else rank + n = t.shape[0] // tp + return _dev(t[rank * n : (rank + 1) * n], d) + + +def _row(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """This rank's input columns (row-parallel input split), even split, on the destination device.""" + n = t.shape[1] // d.tp + return _dev(t[:, d.rank * n : (d.rank + 1) * n], d) + + +def _dev(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """``t`` on the destination device (``d.dev``, set per fill).""" + return t.to(d.dev, non_blocking=True) + + +def _swizzle(scale: torch.Tensor) -> torch.Tensor: + """128x4 swizzle of a row-major ``[N, K / 32]`` UE8M0 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(scale.contiguous()) + + +# --------------------------------------------------------------------------------------------- transforms + + +def _gemma_fold(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A Gemma boundary norm as the plain norm the fused all-reduce epilogue applies: ``w + 1`` in BF16.""" + return t + 1.0 + + +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) + + +def _qkv_sparse_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[Q | K | V | index-Q | index-K]``: Q by attention heads, K/V/index-Q by KV heads, index-K replicated.""" + q, k, v, iq, ik = t + return torch.cat( + [_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0 + ).contiguous() + + +def _qkv_sparse_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v, iq, ik = t + return _swizzle( + torch.cat([_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0) + ) + + +def _qkv_dense_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +def _qkv_dense_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return _swizzle(torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0)) + + +def _index_qk_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[index-Q | index-K]`` in BF16: this rank's index-Q heads and the replicated index-K, the E4M3 values cast + (the scale tensors, the last two sources, are consumed only).""" + iq, ik = t[0], t[1] + return torch.cat([_col(iq, d), _dev(ik, d)], dim=0).to(d.dtype).contiguous() + + +def _gate_up_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return torch.cat([_col(gate, d), _col(up, d)], dim=0).contiguous() + + +def _gate_up_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return _swizzle(torch.cat([_col(gate, d), _col(up, d)], dim=0)) + + +def _row_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _row(t, d).contiguous() + + +def _row_s(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _swizzle(_row(t, d)) + + +def _col_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _col(t, d).contiguous() + + +#: Permute indices of the trtllm-gen expert relayouts, keyed by (shape, family, scale group): one expert's are +#: every expert's, so each is computed once per load. +_PERMUTE_CACHE: dict = {} + + +def _expert_w31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 operand: ``[up ; gate]`` packed-FP4 rows of this rank, then the gated-activation row + reorder and the epilogue-tile shuffle.""" + up, gate = t + w = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + w, _PERMUTE_CACHE, EPILOGUE_TILE_M, is_gated_act_gemm=True + ) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 E4M3 scales: ``[up ; gate]`` rows of this rank, the same row reorder at scale + granularity, then the 128x4 swizzle.""" + up, gate = t + s = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK, is_gated_act_gemm=True + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled.view(torch.uint8).reshape(shape)) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + + +def _expert_w2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 operand: this rank's input columns, then the epilogue-tile row shuffle.""" + w = _row(t, d).contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(w, _PERMUTE_CACHE, EPILOGUE_TILE_M) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 E4M3 scales: this rank's columns, the row shuffle at scale granularity, the swizzle.""" + s = _row(t, d).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w2_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled).view(torch.float8_e4m3fn).reshape(shape) + ) + + +# ------------------------------------------------------------------------------------------------ sources + + +def _one(template: str): + return lambda d, i, t=template: [(_LM + t.format(i=i), None)] + + +def _many(*templates: str): + return lambda d, i, ts=templates: [(tuple(_LM + t.format(i=i) for t in ts), None)] + + +def _per_expert(*suffixes: str): + """One source per routed expert, written at its expert index (all experts are local at tp2).""" + + def build(d, i, sx=suffixes): + out = [] + for e in range(d.num_experts): + keys = tuple(_LM + f"model.layers.{i}.block_sparse_moe.experts.{e}.{s}" for s in sx) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _dense_layers(d): + return range(d.first_sparse) + + +def _sparse_layers(d): + return range(d.first_sparse, d.num_layers) + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense layers, and every layer without the five-way packed + projection (``fuse_qkv_index_projection=False``).""" + return range(d.first_sparse) if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return range(d.first_sparse, d.num_layers) if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else range(d.first_sparse, d.num_layers) + + +_ATTN = "model.layers.{i}.self_attn" +_MOE = "model.layers.{i}.block_sparse_moe" + + +class MiniMaxM3Tp2Weights: + """MiniMax-M3 NVFP4 / sm_107 / tp2: the language model's table and the passes over it. + + ``declare`` allocates every weight it names, ``manifest`` maps checkpoint keys onto them, ``load`` fills + them; ``_targets`` is the one place that decides which parameters exist. + """ + + #: Free the per-expert relayout scratch after each expert stack. + RELEASE_AFTER = "_e_s2" + + WEIGHTS: tuple[W, ...] = ( + # --- the two boundary norms of every layer (Gemma, folded) --- + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + transform=_gemma_fold, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + transform=_gemma_fold, + ), + # --- attention: per-head Q/K norms (Gemma, unfolded), the packed QKV and the output projection --- + W("q_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".q_norm.weight")), + W("k_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".k_norm.weight")), + W( + "qkv_w", + shape=lambda d: (d.qkv_dense_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in "qkv")), + transform=_qkv_dense_w, + layers=_three_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_dense_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many(*(f"{_ATTN}.{p}_proj.weight_scale_inv" for p in "qkv")), + transform=_qkv_dense_s, + layers=_three_way_qkv_layers, + ), + W( + "qkv_w", + shape=lambda d: (d.qkv_sparse_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in ("q", "k", "v", "index_q", "index_k"))), + transform=_qkv_sparse_w, + layers=_five_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_sparse_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + *( + f"{_ATTN}.{p}_proj.weight_scale_inv" + for p in ("q", "k", "v", "index_q", "index_k") + ) + ), + transform=_qkv_sparse_s, + layers=_five_way_qkv_layers, + ), + # The separate BF16 index projection [index-Q | index-K] of the unfused sparse layers. The built-in model + # builds it as an unquantized Linear and fills it through the fused gate/up loader, which casts the + # checkpoint's E4M3 values to BF16 and does not apply their block scales: so does this table (the scale + # tensors are consumed, not applied). + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + src=_many( + _ATTN + ".index_q_proj.weight", + _ATTN + ".index_k_proj.weight", + _ATTN + ".index_q_proj.weight_scale_inv", + _ATTN + ".index_k_proj.weight_scale_inv", + ), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_layers, + ), + W( + "o_w", + shape=lambda d: (d.hidden, d.q_width), + dtype=torch.float8_e4m3fn, + src=_one(_ATTN + ".o_proj.weight"), + transform=_row_w, + ), + W( + "o_s", + shape=lambda d: (d.hidden * (d.q_width // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_ATTN + ".o_proj.weight_scale_inv"), + transform=_row_s, + ), + # --- dense SwiGLU-OAI MLP (layers 0..2) --- + W( + "mlp_gu_w", + shape=lambda d: (2 * d.dense_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight", "model.layers.{i}.mlp.up_proj.weight" + ), + transform=_gate_up_w, + layers=_dense_layers, + ), + W( + "mlp_gu_s", + shape=lambda d: (2 * d.dense_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight_scale_inv", + "model.layers.{i}.mlp.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_dense_layers, + ), + W( + "mlp_dn_w", + shape=lambda d: (d.hidden, d.dense_inter), + dtype=torch.float8_e4m3fn, + src=_one("model.layers.{i}.mlp.down_proj.weight"), + transform=_row_w, + layers=_dense_layers, + ), + W( + "mlp_dn_s", + shape=lambda d: (d.hidden * (d.dense_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one("model.layers.{i}.mlp.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_dense_layers, + ), + # --- MoE (layers 3..59): the fp32 router and its bias, the shared expert --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + dtype=torch.float32, + src=_one(_MOE + ".gate.weight"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "sh_gu_w", + shape=lambda d: (2 * d.shared_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + _MOE + ".shared_experts.gate_proj.weight", _MOE + ".shared_experts.up_proj.weight" + ), + transform=_gate_up_w, + layers=_sparse_layers, + ), + W( + "sh_gu_s", + shape=lambda d: (2 * d.shared_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _MOE + ".shared_experts.gate_proj.weight_scale_inv", + _MOE + ".shared_experts.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_sparse_layers, + ), + W( + "sh_dn_w", + shape=lambda d: (d.hidden, d.shared_inter), + dtype=torch.float8_e4m3fn, + src=_one(_MOE + ".shared_experts.down_proj.weight"), + transform=_row_w, + layers=_sparse_layers, + ), + W( + "sh_dn_s", + shape=lambda d: (d.hidden * (d.shared_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_MOE + ".shared_experts.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_sparse_layers, + ), + # --- the routed experts, NVFP4 (trtllm-gen layouts), all 128 on every rank --- + W( + "e_w31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_per_expert("w3.weight", "w1.weight"), + transform=_expert_w31, + layers=_sparse_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight"), + transform=_expert_w2, + layers=_sparse_layers, + ), + W( + "e_s31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w3.weight_scale", "w1.weight_scale"), + transform=_expert_s31, + layers=_sparse_layers, + ), + W( + "e_s2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w2.weight_scale"), + transform=_expert_s2, + layers=_sparse_layers, + ), + # The per-tensor scalars of every expert, raw; modeling derives the kernels' scalars from them. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + layers=_sparse_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_scale"), + ("isc3", "w3.input_scale"), + ("isc2", "w2.input_scale"), + ("ws2_1", "w1.weight_scale_2"), + ("ws2_3", "w3.weight_scale_2"), + ("ws2_2", "w2.weight_scale_2"), + ) + ), + # --- not per-layer --- + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=_one("model.embed_tokens.weight"), + layers=None, + ), + W( + "final_norm", + shape=lambda d: (d.hidden,), + src=_one("model.norm.weight"), + transform=_gemma_fold, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + sparse = cfg.sparse_attention_config + freq = [int(f) for f in sparse["sparse_attention_freq"]] + tp = mapping.tp_size + head_dim = int(cfg.head_dim) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + index_heads = max(int(sparse["sparse_num_index_heads"]) // tp, 1) + index_dim = int(sparse["sparse_index_dim"]) + return SimpleNamespace( + fused=bool( + getattr( + core.model_config.sparse_attention_config, "fuse_qkv_index_projection", False + ) + ), + num_layers=cfg.num_hidden_layers, + first_sparse=freq.index(1), + hidden=cfg.hidden_size, + head_dim=head_dim, + heads=heads, + kv_heads=kv_heads, + index_heads=index_heads, + index_dim=index_dim, + q_width=heads * head_dim, + qkv_dense_rows=(heads + 2 * kv_heads) * head_dim, + qkv_sparse_rows=(heads + 2 * kv_heads) * head_dim + index_heads * index_dim + index_dim, + dense_inter=int(cfg.dense_intermediate_size) // tp, + shared_inter=int(cfg.shared_intermediate_size) * int(cfg.n_shared_experts) // tp, + moe_inter=int(cfg.intermediate_size) // tp, + num_experts=int(cfg.num_local_experts), + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None)] + return [(f"l{i}_{entry.name}", i) for i in layers] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares (meta under the engine's meta init).""" + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key or key tuple, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, core, weights) -> set: + """The vision tower and its projectors, and the multi-token-prediction head: in the checkpoint, never built + by this target.""" + return {k for k in weights.keys() if k.startswith(_UNBUILT_PREFIXES)} + + def load(self, model, weights) -> None: + """Fill every declared weight and the shell's ``lm_head`` shard; assert coverage both ways.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + assert set(manifest) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest) ^ set(core.w.keys()), + ) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(core, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(core.w[param_key], ckpt, index, transform, weights, d, consumed) + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + # The shell owns the vocab-parallel head: this rank's rows of the full table. + _fill(model.lm_head.weight, _LM + "lm_head.weight", None, _col_w, weights, d, consumed) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + } + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the unbuilt families (vision, MTP head); " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _materialize(entry) -> torch.Tensor: + """``[:]`` on a lazy safetensors slice; a 0-dim tensor (the per-tensor scalars) rejects that index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +def _fill(param, ckpt, index, transform, weights, d, consumed) -> None: + keys = ckpt if isinstance(ckpt, tuple) else (ckpt,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + d.dev = dst.device + # Host tensors: a transform slices this rank's part on the host and moves only that to the device. + src_all = tuple(_materialize(weights[key]) for key in keys) + if transform is not None: + src = transform(src_all if isinstance(ckpt, tuple) else src_all[0], d) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + if src.dim() == 0 and dst.dim() == 0: + pass + elif src.numel() == 1 and dst.numel() == 1: + src = src.reshape(dst.shape) + src = src.to(dst.device, non_blocking=True) + assert tuple(dst.shape) == tuple(src.shape), (ckpt, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + _release(weights, keys) + + +def _release(weights, keys) -> None: + """Drop consumed checkpoint tensors from the engine's dict as the load goes (bounded host memory).""" + if hasattr(weights, "mark_consumed_keys"): + weights.mark_consumed_keys(keys) + else: + for key in keys: + weights.pop(key, None) + + +MODEL_WEIGHTS = MiniMaxM3Tp2Weights() + + +# -------------------------------------------------------------------------------------------- the drafter + + +def _draft_one(template: str): + return lambda d, i, t=template: [(t, None)] + + +def _draft_many(*templates: str): + return lambda d, i, ts=templates: [(tuple(ts), None)] + + +def _draft_fc_norms(d): + return range(d.capture_layers) + + +def _draft_qkv(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +class MiniMaxM3Tp2DraftWeights: + """The one-layer BF16 Eagle3 drafter (``speculative_model``), tp2. + + The fc over the three captured target layers and its per-layer fc_norms, the layer's three norms, its QKV + (input ``[embeds | hidden]``, so twice the hidden width; Q by heads, K/V by KV heads) and o_proj, its gated + MLP, the final norm and the full-vocabulary head (vocab-parallel). The checkpoint's ``embed_tokens`` is the + target's table (the drafter reads the target's); it is a predicted non-load. + """ + + WEIGHTS: tuple[W, ...] = ( + W( + "fc", + shape=lambda d: (d.hidden, d.capture_layers * d.hidden), + src=_draft_one("fc.weight"), + layers=None, + ), + W( + "fc_norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"fc_norm.{i}.weight", None)], + layers=_draft_fc_norms, + ), + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.input_layernorm.weight"), + layers=None, + ), + W( + "hidden_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.hidden_norm.weight"), + layers=None, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.post_attention_layernorm.weight"), + layers=None, + ), + W( + "qkv", + shape=lambda d: (d.qkv_rows, 2 * d.hidden), + src=_draft_many(*(f"layers.0.self_attn.{p}_proj.weight" for p in "qkv")), + transform=_draft_qkv, + layers=None, + ), + W( + "o", + shape=lambda d: (d.hidden, d.q_width), + src=_draft_one("layers.0.self_attn.o_proj.weight"), + transform=_row_w, + layers=None, + ), + W( + "gu", + shape=lambda d: (2 * d.inter, d.hidden), + src=_draft_many("layers.0.mlp.gate_proj.weight", "layers.0.mlp.up_proj.weight"), + transform=_gate_up_w, + layers=None, + ), + W( + "dn", + shape=lambda d: (d.hidden, d.inter), + src=_draft_one("layers.0.mlp.down_proj.weight"), + transform=_row_w, + layers=None, + ), + W("final_norm", shape=lambda d: (d.hidden,), src=_draft_one("norm.weight"), layers=None), + ) + + def dims(self, draft) -> SimpleNamespace: + cfg = draft.draft_config.pretrained_config + mapping = draft.draft_config.mapping + tp = mapping.tp_size + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + return SimpleNamespace( + hidden=cfg.hidden_size, + capture_layers=draft.capture_layers, + head_dim=head_dim, + q_width=heads * head_dim, + qkv_rows=(heads + 2 * kv_heads) * head_dim, + inter=cfg.intermediate_size // tp, + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + if entry.layers is None: + return [(entry.name, None)] + return [(f"{entry.name}{i}", i) for i in entry.layers(d)] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, draft) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + d = self.dims(draft) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, draft, weights) -> set: + """The drafter checkpoint's own embedding: the drafter reads the target's table.""" + return {"embed_tokens.weight"} & set(weights.keys()) + + def load(self, draft, weights) -> None: + d = self.dims(draft) + manifest = self.manifest(draft) + assert set(manifest) == set(draft.w.keys()), set(manifest) ^ set(draft.w.keys()) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(draft, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(draft.w[param_key], ckpt, index, transform, weights, d, consumed) + _fill(draft.lm_head.weight, "lm_head.weight", None, _col_w, weights, d, consumed) + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + f"drafter checkpoint coverage: unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = MiniMaxM3Tp2DraftWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py new file mode 100644 index 000000000000..36a53a834458 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107 / tp4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/__init__.py new file mode 100644 index 000000000000..91d4ff843d5e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107 / tp4 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/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_107__tp4/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_107__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/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_107__tp4/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_107__tp4/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/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_107__tp4/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_107__tp4/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/long_select.py new file mode 100644 index 000000000000..193cc96e72bf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/long_select.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. + +Same contract and bit-identical output as ``torch.ops.trtllm.minimax_m3_select_blocks`` on per-KV-head scores +[num_kv_heads, n_blocks, total_q]: init / local forcing, per-query valid-block count, ties to the smaller block index, +forced -inf picks emitted as -1, ascending block ids with -1 tail padding. + +One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups ``block % 16`` is a lower +bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared memory +and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every candidate +against the whole row (re-read). + +The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp4_select_long`` (this target's copy of ``minimax_m3_select_long`` +with a choice of where it releases the sparse decode), called through this target's ``m3_select_long`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long +from ..ops.attention.m3_select_long_dsl import m3_select_long_dsl + + +def long_select_blocks( + max_score_kv: torch.Tensor, + *, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + dsl: bool = False, +) -> Optional[torch.Tensor]: + """Top-16 select over per-KV-head scores [num_kv_heads, n_blocks, total_q] with 129..9216 blocks. + + Returns the head-major backed [total_q, num_kv_heads, 16] int32 table, or None outside this kernel's shapes + (callers keep the reference op). ``early_trigger`` releases the dependent sparse decode right after the select's + grid wait (``SELECT_EARLY_MAX_ROWS`` in modeling.py); otherwise ``trigger_phase`` 1 releases it once every score + of the row is loaded and 0 at the select's exit (``SELECT_LONG_TRIGGER_PHASE``). ``fast_tail`` ranks rows of at + most 64 candidates on warp 0 (``SELECT_LONG_SL3``; the same picks). ``dsl`` runs the 384-thread kernel + (``SELECT_LONG_DSL``; the same table; ``fast_tail`` does not apply). + """ + if ( + max_score_kv.dtype != torch.float32 + or not max_score_kv.is_cuda + or max_score_kv.dim() != 3 + or not 128 < max_score_kv.shape[1] <= MAX_BLOCKS + or n_valid_blocks.dtype != torch.int32 + or not n_valid_blocks.is_cuda + or n_valid_blocks.shape[0] < max_score_kv.shape[2] + ): + return None + num_kv_heads, _, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, 16), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + return (m3_select_long_dsl if dsl else m3_select_long)( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + fast_tail, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_attn_lc.py new file mode 100644 index 000000000000..8e5027b71ffe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_attn_lc.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 long-context sparse decode attention in one launch. + +Replaces the generic chain's split-K producer (``_gqa_sparse_decode_kernel``, one block per chunk) and the LSE +merge (``_merge_topk_attn_out_kernel``, with its MXFP8 copy for the o_proj) once the top-k block table exists. + +Grid: one cluster of 16 CTAs per query token; CTA ``r`` attends the token's ``r``-th selected block (the Triton +producer's chunk ``r`` at 16 chunks), keeps the partial in shared memory, and after a cluster barrier merges head +``r`` from all 16 partials over DSMEM, writing the BF16 row and, optionally, its MXFP8 copy and swizzled UE8M0 +scales exactly as ``_store_mxfp8_head`` lays them out. + +Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +log-domain recurrence. Accumulation order inside the MMAs and the reductions differs, so results match to rounding, +not bit for bit. + +Schedule: the request length and the query rows (FP8 widened to FP16 in shared memory) are read before the PDL grid +wait, which overlaps them with the select ahead when it releases this grid early (modeling's +``SELECT_EARLY_MAX_ROWS``); +the split-K partials travel by ``st.async`` into each receiving CTA's own mbarrier instead of a cluster-wide barrier. +Neither changes a value or the merge order. + +Geometry: 16 query heads per KV head (TP4), one KV head, head_dim 128, page 128, top-k 16. The kernel is +``torch.ops.trtllm.minimax_m3_sparse_decode_lc``, called through this target's ``m3_sparse_decode_lc`` entry. +""" + +import torch + +from ..ops.attention.m3_sparse_decode_lc import m3_sparse_decode_lc + + +class PrefetchSlot: + """The read-only tensors the next long-context sparse decode prefetches into L2 (the attention layer's o_proj weight + and scales), set by the layer around its decode call and taken by the decode; empty means no prefetch.""" + + def __init__(self) -> None: + self.tensors = None + + def take(self): + tensors, self.tensors = self.tensors, None + return tensors + + +PREFETCH = PrefetchSlot() +# (this layer's id buffer, the next sparse layer's id buffer, its K / V pools, its K / V block scales) +NEXT_KV = PrefetchSlot() + + +def lc_applies(q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output) -> bool: + """Whether the step's operands fit the kernel: TP4 geometry, packed NVFP4 pages, FP8 q, BF16 [T, 16, 128] out.""" + if q.dim() != 3 or tuple(q.shape[1:]) != (16, 128) or q.dtype != torch.float8_e4m3fn: + return False + if q.stride(2) != 1 or q.stride(1) % 16 or q.stride(0) % 16 or q.data_ptr() % 16: + return False + if int(k_paged.shape[1]) != 1 or int(topk_idx.shape[0]) != 1 or int(topk_idx.shape[-1]) != 16: + return False + for t in (k_paged, v_paged): + if ( + t.stride(3) != 1 + or t.stride(2) != 64 + or int(t.shape[3]) != 64 + or t.data_ptr() % 16 + or t.stride(0) % 16 + ): + return False + for t in (k_block_scale, v_block_scale): + if t is None or t.stride(3) != 1 or t.stride(2) != 8 or t.data_ptr() % 8 or t.stride(0) % 8: + return False + return ( + output.dtype == torch.bfloat16 + and output.dim() == 3 + and output.is_contiguous() + and tuple(output.shape) == tuple(q.shape) + ) + + +def m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + linear_softmax=True, + mxfp8=None, + use_pdl=True, + early_trigger=True, + merge_once=False, + prefetch=None, + next_kv=None, + identity=False, + one_div=False, + identity_step=False, +): + """Drop-in for ``minimax_m3_sparse_attn_decode`` on the TP4 NVFP4 shape; callers check ``lc_applies`` first. + + ``topk_idx`` is the head-major [1, total_q, 16] view the Triton producer takes. ``mxfp8`` is ``(q8, sf)`` from + ``MERGE_MXFP8.buffers(2048, dev)`` to also emit the o_proj hand-off. The op is built with the early PDL + trigger only (``early_trigger=True``). ``merge_once`` runs this target's op with the merge weights formed once + per CTA (bitwise the same output); with it, ``prefetch`` (up to two contiguous tensors) is pulled into L2 by the + op's CTAs before their grid dependency. + """ + total_q = int(q.shape[0]) + q8, sf = ( + (mxfp8[0].view(torch.uint8), mxfp8[1].view(torch.uint8)) + if mxfp8 is not None + else (None, None) + ) + m3_sparse_decode_lc( + q.view(torch.uint8), + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + output.view(total_q, 2048), + q8, + sf, + 2048 // 128, + int(decode_query_len), + float(sm_scale), + bool(linear_softmax), + bool(use_pdl), + bool(early_trigger), + merge_once=bool(merge_once), + identity=bool(identity) and bool(merge_once), + identity_step=bool(identity_step) and bool(merge_once), + one_div=bool(one_div) and bool(merge_once), + prefetch=prefetch if merge_once else None, + next_kv=( + tuple(next_kv[:2]) + tuple(t.view(torch.uint8) for t in next_kv[2:]) + if merge_once and next_kv is not None + else None + ), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_lowm.py new file mode 100644 index 000000000000..00c6c6887575 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_lowm.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M kernels (``DRAFT_LOWM``). + +The drafter's bias-free BF16 projections (qkv, o_proj, gate_up, down, the fc; not the head) call ``draft_linear`` +first: a covered row count runs ``ops/gemm/m3_draft_lowm`` at the tactic of its (N, K) table, every other call the +plain GEMM. ``prebuild`` compiles every covered call before the executor's warmup; a forward never compiles one, and a +covered call it finds uncompiled is a miss. +""" + +from __future__ import annotations + +from typing import FrozenSet, Iterable, Optional, Set, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm.logger import logger + +from ..ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for +from . import m3_report + +PATH = "m3_draft_lowm" + + +# ((N, K), M) of every call the target covers (its DRAFT_LOWM keys), set by prebuild. +_KEYS: FrozenSet[Tuple[Tuple[int, int], int]] = frozenset() +# (N, K, M, PDL) of every call prebuild compiled. +_BUILT: Set[Tuple[int, int, int, bool]] = set() +# (N, K, M) of every call that logged its engagement line. +_ENGAGED: Set[Tuple[int, int, int]] = set() +# DRAFT_AR_PUSH: the weight whose next low-M split-K call also pushes its rows into a boundary all-reduce +# (data pointer, the bucket's multicast words, its flags, rank, world, max tokens), set by arm_push; and whether the +# last armed call pushed: the data pointer of the output it pushed, 0 if none (read once by take_pushed). Host state +# set and read in one forward's issue order. +_PUSH_ARM: Optional[tuple] = None +_PUSHED = 0 + + +def arm_push(weight: torch.Tensor, workspace, rows: int) -> None: + """Arm ``weight``'s next covered split-K call to push into ``workspace``'s bucket for ``rows`` (see _PUSH_ARM).""" + global _PUSH_ARM, _PUSHED + from . import m3_mc_allreduce + + bucket = m3_mc_allreduce._bucket(workspace, rows) + flags = None if bucket is None else bucket.flags_for(8) + _PUSHED = 0 + _PUSH_ARM = None + if bucket is not None and flags is not None and bucket.max_tokens == 8: + _PUSH_ARM = ( + weight.data_ptr(), + bucket.mc.data_ptr(), + flags.data_ptr(), + bucket.rank, + bucket.world_size, + bucket.max_tokens, + ) + + +def take_pushed(x: torch.Tensor) -> bool: + """Whether the armed call pushed its rows; disarms. ``x`` is what the boundary all-reduce is about to reduce: it + must be the pushing call's own output (the prepushed all-reduce reads only the pushed words, never ``x``), so an + operation between that call and the all-reduce raises instead of silently reducing the wrong rows.""" + global _PUSH_ARM, _PUSHED + pushed, _PUSHED, _PUSH_ARM = _PUSHED, 0, None + if pushed and x.data_ptr() != pushed: + raise RuntimeError( + "the drafter's pushed rows are not the boundary all-reduce's input (an operation ran in between)" + ) + return bool(pushed) + + +def draft_linear( + weight: torch.Tensor, x: torch.Tensor, bias: Optional[torch.Tensor] +) -> Optional[torch.Tensor]: + """x [M, K] @ weight^T through the low-M kernel when the target covers it (``prebuild``'s keys); None for the + plain GEMM. The split-K ('s') tactics run this target's copy of FlashInfer's kernel (kernels/m3_lowm_splitk, bit + for bit the same output at every certified cell), which also carries the drafter-boundary push epilogue.""" + if bias is not None or x.dim() != 2: + return None + n, k = weight.shape + m = int(x.shape[0]) + if ((n, k), m) not in _KEYS: + m3_report.uncovered(PATH, f"the drafter's {n}x{k} linear past its tactic table's rows") + return None + if not x.is_contiguous(): + # The first pass's fc may read a column window of the captured hidden states; the plain GEMM takes strides. + m3_report.uncovered(PATH, "a strided activation (the plain GEMM reads it in place)") + return None + use_pdl = get_env_enable_pdl() + if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + m3_report.missed(PATH, m, "the activation or weight is not BF16") + return None + if (n, k, m, use_pdl) not in _BUILT: + m3_report.missed(PATH, m, f"the {n}x{k} kernel at {m} rows was not compiled before warmup") + return None + kind, tactic = tactic_for(n, k, m) + if (n, k, m) not in _ENGAGED: + _ENGAGED.add((n, k, m)) + logger.info( + f"[M3 sm_107 tp4] fast path engaged: m3_draft_lowm {n}x{k} at {m} rows (device {x.device.index})" + ) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + global _PUSHED + ar = None + if _PUSH_ARM is not None and _PUSH_ARM[0] == weight.data_ptr() and kind == "s": + ar = _PUSH_ARM[1:] + _PUSHED = out.data_ptr() + return m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, True, ar) + + +def installed(weight: torch.Tensor) -> bool: + """Whether the target covers ``weight`` (bias-free BF16 [N, K]) at some row count (``prebuild``'s keys).""" + if weight.dtype != torch.bfloat16 or weight.dim() != 2: + return False + return any(nk == tuple(weight.shape) for nk, _ in _KEYS) + + +def prebuild(weights: Iterable[torch.Tensor], keys: Iterable[Tuple[Tuple[int, int], int]]) -> int: + """Take the target's covered ((N, K), M) ``keys``, then compile (and run once) every covered row count of every + covered drafter weight, outside capture, on the current stream. Returns the number of calls compiled.""" + global _KEYS + _KEYS = frozenset(keys) + use_pdl = get_env_enable_pdl() + built = 0 + shapes = {tuple(weight.shape): weight for weight in weights if installed(weight)} + for (n, k), weight in sorted(shapes.items()): + for nk, m in sorted(key for key in _KEYS if key[0] == (n, k)): + if (n, k, m, use_pdl) in _BUILT: + continue + kind, tactic = tactic_for(n, k, m) + x = torch.zeros(m, k, dtype=torch.bfloat16, device=weight.device) + out = torch.empty(m, n, dtype=torch.bfloat16, device=weight.device) + m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, True) + if kind == "s": + # The push variant a boundary all-reduce may arm (o_proj / down): compiled only, never run here + # (its pushes need the all-reduce's exchange and a matching call). + from .m3_lowm_splitk import compile_push + + compile_push(x, weight.t(), out, use_pdl, tactic) + _BUILT.add((n, k, m, use_pdl)) + built += 1 + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_step.py new file mode 100644 index 000000000000..45a64dfe9210 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_step.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 draft pass for the MiniMax-M3 drafter's vocab-sharded BF16 head (TP decode, <= 8 requests). + +One kernel replaces a draft pass's head and the draft loop's per-pass bookkeeping, at the in-tree precision: + +* The drafter's BF16 lm_head shard is kept a second time, built by ``prepare_head`` before the executor's warmup: in + the lossless 12-bit coding of the target head (``ops.gemm.m3_target_head``; 75% of the bytes, kept only when it + decodes to the BF16 shard bit for bit), else in ``mma.m16n8k16`` A-fragment order (``bf16_fragments``, the same + values). The kernel streams it with BF16 MMAs and FP32 accumulation over the BF16 head input rows; a coded + exception decodes to +0 and its exact product x * w is added at its tile's flush. +* Each logit is rounded to BF16 (the in-tree logits are the BF16 head GEMM's output) and keyed + ``ordered(value) << 32 | ~global_index`` with -0 as +0, so the larger value wins and ties take the lowest index, as + the in-tree ``torch.max`` / ``torch.argmax`` over the gathered logits (``greedy_sample_draft_with_tp_gather``). +* The last CTA exchanges each row's key across the TP ranks through a multicast mapping (Lamport slots, see + ``m3_shard_accept``), then stores the draft token, the next pass's input id and position, and advances the draft KV + lengths as the in-tree draft loop does between passes. + +The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu) and called through the +catalog entry ``ops.spec.m3_draft_step``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok +from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments, max_tile_pieces +from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op +from . import m3_report +from .m3_shard_accept import AcceptWorkspace, create_workspace + +# The draft step launches this many CTAs instead of one per SM (3126 tiles: 200 CTAs x 15 + 9 x 14, at most 2 warp +# pieces per tile), so the CTAs own near-equal whole-tile counts and the stream does not end on the tiles' straggler +# tail (-1.04 / -1.18 us per pass at 1 / 4 rows vs 212 on Rubin TP4). A different static warp segmentation: +# exact-class. +GRID = 209 # capacity: fixed: 3126 head tiles of the TP4 shard, at most 2 warp pieces per tile on 209 CTAs +# The draft head kernel sums a tile that several warps stream in warp order (fixed by the launch shape) instead of by +# shared-memory atomics in arrival order: with fewer than 16 tiles per CTA a tile can take 3-4 pieces and the atomics' +# order changed the fp32 sum run to run; with 2 pieces the sum is the atomics' bit for bit. +FIXED_ORDER_PATH = "m3_draft_step_fixed_order" + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the 12-bit coded stream + with its exception lists (``coded``), or the fragment-ordered BF16 copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + coded: bool = False + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + exp_base: int = 0 + + +def head_supported(lm_head) -> bool: + """A bias-free BF16 head whose local shard the kernel can own: whole 16-row tiles, K a multiple of 16, at most + MAX_TILES_PER_CTA tiles per CTA of a one-CTA-per-SM grid.""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + if getattr(lm_head, "bias", None) is not None: + return False + n, k = weight.shape + if n % 16 or k % 16: + return False + return (n // 16 + _grid() - 1) // _grid() <= MAX_TILES_PER_CTA + + +def prepare_head(lm_head, tp_rank: int, coded: bool) -> DraftHead: + """The streamed copy of ``lm_head``'s local shard (global ids from tp_rank * rows): its 12-bit coding when + ``coded``, K is a multiple of 64 and the coding decodes to the shard bit for bit, else (fail + closed) the fragment-ordered BF16 copy. + Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_draft_head``).""" + weight = lm_head.weight + offset = tp_rank * weight.shape[0] + with torch.no_grad(): + w = weight.detach() + if coded and w.shape[1] % 64 == 0: + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return DraftHead( + weight=weight, + wf=packed, + vocab_offset=offset, + coded=True, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + exp_base=base, + ) + del packed, exc_off, exc_pos, exc_w + wf = bf16_fragments(w) + return DraftHead(weight=weight, wf=wf, vocab_offset=offset) + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _grid() -> int: + """The draft step's CTA count: ``GRID`` (at most one per SM).""" + return min(GRID, _num_sms()) + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace ([2][MAX_TOKENS][world] slots), created collectively on the first eager call; None + during a CUDA-graph capture that finds it missing.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = create_workspace(mapping, device, MAX_TOKENS) + return _WS + + +def _fixed_order( + rows: int, head: DraftHead, x: torch.Tensor, num_ctas: int, split: bool = False +) -> bool: + """The fixed warp-order tile sum where the launch lets a tile take more than 2 warp pieces (max_tile_pieces; a + split or equal-chunk launch always), with its engagement line; False keeps the atomics, deterministic on this + shape.""" + eq = getattr(head, "bnd", None) is not None + if not (split or eq): + chunks = x.shape[1] // (16 if head.exc_off is None else 64) + if ( + max_tile_pieces( + head.wf.shape[0] if head.exc_off is None else head.exc_off.numel() - 1, + num_ctas, + chunks, + ) + <= 2 + ): + return False + m3_report.taken( + FIXED_ORDER_PATH, + rows, + "the draft head sums shared tiles in warp order", + ) + return True + + +def draft_step( + head: DraftHead, + x: torch.Tensor, + rows: int, + x_rows: Optional[torch.Tensor], + mapping, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, +) -> bool: + """One fused greedy draft pass over ``rows`` requests (see the module docstring); False when the exchange + workspace is missing during a capture. Collective over the TP group. + + Args: + head: the prepared head of this rank (``prepare_head``), coded or BF16 fragments. + x: [*, K] BF16 head input rows (the drafter's final-norm output); row ``x_rows[b]`` (int64) or ``b`` is + request b's. + draft_out: int32 [rows, >= step + 1], unit column stride: column ``step`` receives the tokens. + ids_next / pos_next: int32 [rows]: the tokens, and ``pos_src[x_rows[b] or b] + 1`` (pos_src int32). + kv_lens: int32 draft KV lengths: ``-= draft_len - num_acc`` with ``seq_lens = 1`` on pass 0 when ``num_acc`` + is given, ``+= 1`` on later passes. + """ + ws = workspace(mapping, x.device) + if ws is None: + return False + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + _grid(), + True, + None, + head.exc_off, + head.exc_pos, + head.exc_w, + head.exp_base, + None, # no fold tail: the forward's tail runs the standalone one + fixed_order=_fixed_order(rows, head, x, _grid()), + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_drafter_attn.py new file mode 100644 index 000000000000..e5fa1bb5e945 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_drafter_attn.py @@ -0,0 +1,405 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's decode attention on Rubin: ``M3DrafterFullCtxFmha``. + +The drafter layer is a plain ``TrtllmAttention`` over its own FP8 paged cache (the full context, every draft pass). +On SM 10.7 its generation phase runs ``thop.attention``'s in-tree trtllm-gen runner, whose split is capped at 8 CTAs +per (request, KV head): at TP4 (1 KV head per rank) 45.4 us per pass at 165k tokens on 212 SMs. + +This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append +into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context +PrimTS split-KV decode of ``kernels/m3_fullctx_decode`` at the same pinned recipes as the dense layers. + +With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead +(``minimax_m3_sm107tp4_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page +pool and block table taken from the KV cache pool pointers and mapping the preprocess reads. + +It serves generation-only steps of the covered (query rows per request, requests) shapes; the target installs it +first in the drafter layer's FMHA list (``install``), so every other step falls through to the stock libraries. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + FlashInferTrtllmGenFmha, +) +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode + +from . import m3_fullctx_decode as fullctx + +# With PDL, the fused RoPE launch takes the maximum shared-memory carveout so the drafter decode's large-smem CTAs +# co-reside with it (timing only). +ROPE_CARVEOUT = True + + +def fwd_has_no_mrope(fwd) -> bool: + return getattr(fwd, "mrope_position_deltas", None) is None + + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + +class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): + """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + + def __init__( + self, + attn: "TrtllmAttention", + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + pdl: bool = False, + on_rope_taken=None, + ) -> None: + super().__init__(attn) + self._recipes = recipes + self._max_kv_len = max_kv_len + self._own_reducer = own_reducer + self._fused_rope_mode = fused_rope_mode + self._pdl = pdl + self._on_rope_taken = on_rope_taken + self._on_taken = on_taken + self._on_uncovered = on_uncovered + self._on_missed = on_missed + + def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: + nreq = int(metadata.num_generations) + if nreq <= 0 or int(q.shape[0]) % nreq: + return None + return int(q.shape[0]) // nreq, nreq + + def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: + if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: + return False + if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: + return False + attn = self.attn + if ( + attn.num_heads != fullctx.HEADS + or attn.num_kv_heads != fullctx.KV_HEADS + or attn.head_dim != fullctx.HEAD_DIM + or not QuantMode(attn.quant_mode).has_fp8_kv_cache() + ): + self._on_uncovered("not the TP4 FP8 drafter geometry") + return False + # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's + # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. + if ( + forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant + or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig + ): + self._on_uncovered("per-step KV cache scales") + return False + if int(metadata.max_seq_len) > self._max_kv_len: + self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") + return False + shape = self._shape(q, metadata) + if shape is None or shape not in self._recipes: + self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") + return False + # Compiled before warmup (``install``); the forward never compiles one. + if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: + self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") + return False + return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) + + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + attn = params.attn + width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim + rotary_dim = int(attn.rope_params.dim) + return ( + self._fused_rope_mode is not None + and params.qkv_input is not None + and attention_input.dim() == 2 + and attention_input.dtype == torch.bfloat16 + and attention_input.is_contiguous() + and int(attention_input.shape[1]) == width + and not params.is_cross + and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) + and attn.rotary_cos_sin is not None + and attn.rotary_cos_sin.dtype == torch.float32 + and fwd_has_no_mrope(params.fwd) + and 0 < rotary_dim <= attn.head_dim + and rotary_dim % 16 == 0 + and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) + # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo + # runs for several requests only) and reads the workspace's, so it keeps that path. + and not (self._multi_token(params) and int(params.batch_size) == 1) + ) + + @staticmethod + def _multi_token(params: FmhaParams) -> bool: + return ( + params.spec_decoding_generation_lengths is not None + and int(params.attn.predicted_tokens_per_seq) > 1 + ) + + def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" + if not self._fused_rope_applies(params, attention_input): + return False + attn = params.attn + meta = params.meta + qlen, nreq = int(params.input_seq_length), int(params.batch_size) + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + if plan is None: + return False + # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary + # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the + # block table is the pool's [requests, 2, blocks] K / V sub-page indices. + pool_index, layer_in_pool = ( + int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] + ) + tpb = int(params.tokens_per_block) + sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim + base = ( + int(meta.host_kv_cache_pool_pointers[pool_index, 0]) + + layer_in_pool * int(params.kv_factor) * sub_page_bytes + ) + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) + ) + tables = meta.kv_cache_block_offsets[ + pool_index, params.seq_offset : params.seq_offset + nreq + ] + multi = self._multi_token(params) + if not torch.cuda.is_current_stream_capturing(): + # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess + # path checks; the op itself writes V at the table's V sub-page. + tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(tables.shape[-1], device=tables.device) + in_use = cols[None, :] < tpb_held[:, None] + offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + q = torch.empty( + int(params.num_tokens), + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=attention_input.device, + ) + + torch.ops.trtllm.minimax_m3_sm107tp4_drafter_rope_fp8_kv( + attention_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + int(attn.rope_params.dim), + tpb, + self._fused_rope_mode, + self._pdl, + ROPE_CARVEOUT, + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + tables[:, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + if self._on_rope_taken is not None: + self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") + self._on_taken(params.num_tokens) + return True + + def run_generation(self, params: FmhaParams) -> None: + attention_input = params.qkv_input if params.qkv_input is not None else params.query_input + if attention_input is None: + raise RuntimeError("the drafter decode requires QKV or query input.") + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + if self._run_fused_rope(params, attention_input): + return + # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 + # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). + fp8_context_fmha = False + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + _fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + _window_left, + is_multi_token_gen, + ) = thop.trtllm_gen_generation_preprocess( + attention_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + fp8_context_fmha, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + ) + qlen, nreq = params.input_seq_length, params.batch_size + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the + # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row + # pass is never multi-token. + if plan is None or (is_multi_token_gen and qlen == 1): + raise RuntimeError( + f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" + ) + rows = params.num_tokens * attn.num_heads * attn.head_dim + q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) + q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) + # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 + # of the pool (the two views have the same page count, as the kernel requires). + if not torch.cuda.is_current_stream_capturing(): + # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. + lengths = params.spec_decoding_generation_lengths + if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): + raise RuntimeError( + f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" + ) + # Only the blocks a request holds are laid out; the table's tail entries are unset. + tpb = int(params.tokens_per_block) + held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(block_tables.shape[-1], device=block_tables.device) + in_use = cols[None, :] < held[:, None] + offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. + # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + block_tables[:nreq, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + self._on_taken(params.num_tokens) + + +def install( + draft_model: torch.nn.Module, + recipes, + max_kv_len: int, + own_reducer: bool, + pdl: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + on_rope_taken=None, +) -> int: + """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention + backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it + was installed on.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + count = 0 + device = next(draft_model.parameters()).device + # The backend is not a module: each attention module of the drafter holds it as ``attn``. + for module in draft_model.modules(): + backend = getattr(module, "attn", None) + if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): + continue + if not bool((backend.kv_cache_scaling_factor == 1.0).all()): + on_uncovered("a scaled FP8 KV cache") + continue + fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) + libs = backend._fmha_manager.fmha_libs + if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): + libs.insert( + 0, + M3DrafterFullCtxFmha( + backend, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode=fused_rope_mode, + pdl=pdl, + on_rope_taken=on_rope_taken, + ), + ) + count += 1 + return count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fold_tail.py new file mode 100644 index 000000000000..79aff5cb80ce --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fold_tail.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The speculative step's store and next-step staging at the tail of the forward (one Triton launch). + +For each of the first ``n`` rows (``n`` read from device memory, so a captured graph whose step was not armed writes +nothing), with ``s = slots[i]``: + +* the spec sampler's per-slot store, exactly as ``m3_spec_tail.slot_scatter`` writes it (zero-padded to each store + width): ``new[s][c]``, ``next_new[s][c]``, ``draft[s][c]`` and ``lens[s]``; +* the engine's next-step inputs, exactly as ``m3_spec_tail.step_gather`` writes them from that store for a step whose + rows are these rows in this order (no new tokens, first previous row 0): ``ids[i*T + t] = next_new[s][t]``, + ``draft_ids[i*D + d] = draft[s][d]``, ``pos_off[i*T + t] = lens[s]``, ``kv_off[i] = lens[s] - T``. + +Integer copies only: the values are the ones the sampler's and the engine's own copies write. + +With ``cpp`` the same writes run as this target's C++ op (trtllm::minimax_m3_sm107_tp4_fold_tail, +minimaxM3Sm107Tp4/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fold_tail_kernel( + ctrl, + slots, + src_new, + sn_s0, + sn_s1, + SN_W: tl.constexpr, + dst_new, + dn_s0, + dn_s1, + DN_W: tl.constexpr, + src_next, + sx_s0, + sx_s1, + SX_W: tl.constexpr, + dst_next, + dx_s0, + dx_s1, + DX_W: tl.constexpr, + src_draft, + sd_s0, + sd_s1, + SD_W: tl.constexpr, + dst_draft, + dd_s0, + dd_s1, + DD_W: tl.constexpr, + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + T: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(ctrl) + if i < n: + s = tl.load(slots + i).to(tl.int64) + c = tl.arange(0, BLOCK) + new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) + tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) + nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) + tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) + drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) + tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) + ln = tl.load(src_lens + i) + tl.store(dst_lens + s, ln) + # The next step's inputs read the store's padded rows: zero past each source width. + tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) + tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) + tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) + tl.store(kv_off + i, ln - T) + + +def fold_tail( + ctrl: torch.Tensor, + slots: torch.Tensor, + rows: int, + src_new: torch.Tensor, + dst_new: torch.Tensor, + src_next: torch.Tensor, + dst_next: torch.Tensor, + src_draft: torch.Tensor, + dst_draft: torch.Tensor, + src_lens: torch.Tensor, + dst_lens: torch.Tensor, + ids: torch.Tensor, + draft_ids: torch.Tensor, + pos_off: torch.Tensor, + kv_off: torch.Tensor, + tokens_per_step: int, + cpp: bool = False, +) -> bool: + """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs + are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width].""" + two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) + one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) + if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): + return False + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) + or slots.dtype != torch.long + ): + return False + t_step = tokens_per_step + d_step = t_step - 1 + if ( + rows <= 0 + or slots.numel() < rows + or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) + or ids.numel() < rows * t_step + or pos_off.numel() < rows * t_step + or draft_ids.numel() < rows * d_step + or kv_off.numel() < rows + ): + return False + widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) + widths += (src_draft.shape[1], dst_draft.shape[1], t_step) + if cpp and max(widths) <= 32: + torch.ops.trtllm.minimax_m3_sm107_tp4_fold_tail( + ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, + src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, + ) # fmt: skip + return True + block = triton.next_power_of_2(max(widths)) + _fold_tail_kernel[(rows,)]( + ctrl, + slots, + src_new, + src_new.stride(0), + src_new.stride(1), + src_new.shape[1], + dst_new, + dst_new.stride(0), + dst_new.stride(1), + dst_new.shape[1], + src_next, + src_next.stride(0), + src_next.stride(1), + src_next.shape[1], + dst_next, + dst_next.stride(0), + dst_next.stride(1), + dst_next.shape[1], + src_draft, + src_draft.stride(0), + src_draft.stride(1), + src_draft.shape[1], + dst_draft, + dst_draft.stride(0), + dst_draft.stride(1), + dst_draft.shape[1], + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + t_step, + d_step, + block, + num_warps=1, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_force_vec.py new file mode 100644 index 000000000000..be04f26f9aff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_force_vec.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's +``force_num_accepted_tokens``), in one launch. + +Triton: ``force_vec_kernel``: one CTA bumps the shared step counter and draws each request's count from the fixed-seed +pool with the in-tree (counter, slot) hash, exactly as ``_apply_force_accepted_tokens`` does. ``force_vec_req_kernel`` +(``FORCE_AL_PER_REQUEST``): each request draws from the same pool at its own key, its KV length at this verify. +Reached through the ``ops/spec`` entries ``m3_force_vec`` and ``m3_force_vec_req``. + +One compiled variant for every batch and accept count: specializing them (== 1, divisible by 16) would compile +mid-serving the first time a step reaches a new one, past the warm-up. +""" + +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_kernel( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK: tl.constexpr, +): + c = tl.load(counter) + 1 + slots = tl.arange(0, BLOCK) + tl.store(counter, c) + idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask + live = slots < batch_size + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) + + +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_req_kernel( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK: tl.constexpr, +): + # One draw per request from its own key (its KV length at this verify): no shared counter, so a request's forced + # counts depend only on its own prompt and progress, not on what else the server ran before or beside it. + slots = tl.arange(0, BLOCK) + live = slots < batch_size + k = tl.load(keys + slots, mask=live, other=0).to(tl.int64) + idx = (k * key_stride) & pool_mask + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fullctx_decode.py new file mode 100644 index 000000000000..ef6b9ee0d1ad --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fullctx_decode.py @@ -0,0 +1,313 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 full-context FP8 decode attention on Rubin: ``fullctx_decode``. + +The dense layers (0-2) attend every token of their FP8 E4M3 (unscaled) paged KV cache. On SM 10.7 the stock call +(flashinfer trtllm-gen, ``minimax_m3_trtllm_gen_dense_decode``) caps its split at 16 cluster CTAs per (request, KV +head): at TP4 (16 query heads and 1 KV head per rank) one request runs on 16 CTAs of 212 and takes 33.4-33.8 us +per call at 165k tokens. + +This runs the same arithmetic -- FP8 Q (the same E4M3 rounding of q), FP8 K/V, P quantized to E4M3 per split -- on +the vendored PrimTS decode kernel (``attention/backends/prims_ts``, SM 10.7 through the sm_100f family) at a pinned +split-KV recipe, which fans one request out over 96-108 CTAs. The split partials are merged by +``torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce`` (``cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/ +fullCtxReduce.cu``, head-count generic, reused by name): unnormalized fp32 accumulation over the active splits, one +final division, one BF16 rounding, the merge form of stock trtllm-gen. Only the split changes against stock +(``modeling.FULLCTX_DECODE_RECIPES``). + +A plan is compiled per (device, query rows per request, requests) by ``prebuild`` before the executor's warmup; the +forward only reads plans (``plan_for``). The plan is built from PrimTS's own config, compile and workspace helpers at +an explicit recipe; the PrimTS wrapper's planning policy is not used, so no PrimTS state is replaced. With +``own_reducer`` off the PrimTS reducer merges into FP16 and the result is cast to BF16. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Mapping, Optional, Tuple + +import torch + +# The TP4 per-rank geometry: 16 query heads over 1 KV head. +HEADS = 16 +KV_HEADS = 1 # capacity: fixed: the tp4 KV heads per rank +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim +PAGE = 128 +# The per-request K/V length PrimTS's config selection sizes its policy for (Q/KV tiles, launch mode); the pinned +# recipes were measured with it. The compiled capacity (``max_kv_len``) may be larger. +POLICY_KV_LEN = ( + 1 << 20 +) # capacity: fixed: the PrimTS policy's K/V length bound (the plans' compile shape) +HEADS_PER_KV = HEADS // KV_HEADS # capacity: fixed: the tp4 query heads per KV head + +_plans: Dict[Tuple[int, int, int, bool], "FullCtxPlan"] = {} +_lock = threading.Lock() + + +class FullCtxPlan: + """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" + + def __init__( + self, + device: torch.device, + qlen: int, + nreq: int, + max_kv_len: int, + recipe: Mapping, + own_reducer: bool, + pdl: bool, + ): + from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec + from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( + get_max_active_clusters_for_cluster_size, + make_decode_config, + ) + + self._tsdec = tsdec + self.qlen = qlen + self.nreq = nreq + self.pdl = pdl + self.recipe = dict(recipe) + splits = int(self.recipe["splits"]) + fp8 = tsdec._cutlass_dtype("float8_e4m3fn") + # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer + # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents + # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an + # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the + # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own + # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge + # (minimax_m3_sm107tp4_fullctx_reduce) waits as its first instruction. + cfg = make_decode_config( + headdim=HEAD_DIM, + args={"use_pdl": pdl, **self.recipe.get("config", {})}, + seq_len_q=qlen, + seq_len_kv=min(max_kv_len, POLICY_KV_LEN), + batch_size=nreq, + num_heads_q=HEADS, + num_heads_kv=KV_HEADS, + q_dtype=fp8, + k_dtype=fp8, + v_dtype=fp8, + o_dtype=tsdec._cutlass_dtype("float16"), + qkv_layout="pagedKv", + num_tokens_per_page=PAGE, + storage_tokens_per_page=PAGE, + split_kv_mode="gmem_reduction_with_separate_kernel", + splits_kv=splits, + max_splits_kv=splits, + min_loop_iters_per_split=1, + mask_type="causal" if qlen > 1 else "dense", + auto_tuner=True, + split_kv=True, + ) + _check_layout(cfg, qlen, splits) + with torch.cuda.device(device): + max_active = get_max_active_clusters_for_cluster_size(1) + spec = tsdec._decode_launch_spec_from_config( + cfg, + batch_size=nreq, + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + seq_len_q=qlen, + max_active_clusters=int(max_active), + ) + compile_spec = tsdec._make_decode_compile_spec( + spec, + device_index=_index(device), + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + page_size=PAGE, + max_kv_len=max_kv_len, + seq_len_q=qlen, + q_dtype_key="float8_e4m3fn", + k_dtype_key="float8_e4m3fn", + v_dtype_key="float8_e4m3fn", + output_dtype_key="float16", + use_packed_q=False, + kv_prefix_mode="dynamic", + kv_lengths_mode="dynamic", + ) + self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) + layout = tsdec._make_decode_workspace_layout( + spec.scratch_shapes, + torch.float16, + use_separate_reduction_kernel=True, + use_split_kv=True, + ) + self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) + self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) + self.workspace.split_kv_counter.zero_() + self.workspace.cu_seqlens_q.zero_() + self.workspace.attention_sinks.zero_() + rows = HEADS_PER_KV * qlen + if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): + raise RuntimeError( + f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " + f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" + ) + self.own_reducer = own_reducer + self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) + self.tile_kv = int(cfg.tile_size_kv) + self.insts_kv = int(cfg.num_insts_kv) + self.causal = qlen > 1 + shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) + self.staging = torch.empty(shape, dtype=torch.float16, device=device) + + def run( + self, + q: torch.Tensor, + k_pages: torch.Tensor, + v_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + mxfp8: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> bool: + """Decodes into output; with mxfp8 (E4M3 rows, UE8M0 scales) and a merge op that takes them, the merge + also writes the MXFP8 copy of output. Returns whether it did.""" + # The kernel takes a compact query; a column view of the fused QKV projection is copied once. + q = q.reshape(self.staging.shape).contiguous() + runtime = self._tsdec._prepare_decode_runtime_unchecked( + q, + (k_pages, v_pages), + None, + output_dtype=torch.float16, + bmm1_scale=sm_scale, + bmm2_scale=1.0, + out=self.staging, + ) + ws = self.workspace + self._main( + runtime.q, + runtime.k_cache, + runtime.v_cache, + runtime.k_sf_cache, + runtime.v_sf_cache, + runtime.out, + seq_lens, + ws.cu_seqlens_q, + block_table, + block_table, # Dense decode constexpr-elides the membership table. + ws.partial_o, + ws.partial_stats, + ws.split_kv_counter, + ws.attention_sinks, + runtime.num_physical_pages, + runtime.k_page_stride, + runtime.k_head_stride, + runtime.k_token_stride, + runtime.v_page_stride, + runtime.v_head_stride, + runtime.v_token_stride, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + if self.own_reducer: + args = ( + ws.partial_o, + ws.partial_stats, + seq_lens, + output, + self.qlen, + self.q_tokens_per_cta, + self.tile_kv, + self.insts_kv, + self.causal, + self.pdl, + ) + if mxfp8 is not None and emits_mxfp8(): + torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce(*args, mxfp8[0], mxfp8[1]) + return True + torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce(*args) + return False + self._primts_reducer( + runtime.out, + seq_lens, + ws.cu_seqlens_q, + ws.partial_o, + ws.partial_stats, + ws.attention_sinks, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + output.view(self.staging.shape).copy_(self.staging) + return False + + +def _check_layout(cfg, qlen: int, splits: int) -> None: + """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. + + Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is + ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding + window, attention sinks or packed Q. + """ + problems = [] + if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): + problems.append("not split-KV with the separate GMEM reducer") + if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: + problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") + if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: + problems.append("cluster-SMEM reduction or a persistent scheduler") + if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: + problems.append("packed Q, a sliding window or attention sinks") + if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): + problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") + if qlen > 1: + if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): + problems.append("multi-token Q without the grouped [token, head] tile layout") + elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): + problems.append( + f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" + ) + if problems: + raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + + +def emits_mxfp8() -> bool: + """Whether the loaded merge op can also write the o_proj MXFP8 hand-off (its schema takes q8 / sf).""" + op = getattr(torch.ops.trtllm, "minimax_m3_sm107tp4_fullctx_reduce", None) + return op is not None and "q8" in str(op.default._schema) + + +def has_own_reducer() -> bool: + """Whether this build carries the BF16 split-KV merge op.""" + return hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_fullctx_reduce") + + +def plan_for( + device: torch.device, qlen: int, nreq: int, own_reducer: bool = True +) -> Optional[FullCtxPlan]: + """The prebuilt plan for this shape, or None.""" + return _plans.get((_index(device), qlen, nreq, own_reducer)) + + +def prebuild( + device: torch.device, + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + pdl: bool, +) -> int: + """Compile the plan of every (query rows per request, requests) in ``recipes``. Call before warmup.""" + built = 0 + for (qlen, nreq), recipe in recipes.items(): + key = (_index(device), qlen, nreq, own_reducer) + with _lock: + if key not in _plans: + _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) + built += 1 + return built + + +def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: + """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" + return pool[0::subpages_per_slot], pool[1::subpages_per_slot] + + +def _index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score.py new file mode 100644 index 000000000000..fad75d44694a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score.py @@ -0,0 +1,638 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with two schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer``) through ``IndexScoreRunner`` +and called through this target's ``m3_index_score`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, griddepcontrol_wait + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # No explicit launch_dependents: the only consumer is the block selector, + # which reduces over every block this grid writes, so triggering it early + # would leave it resident and spinning in its own wait. The implicit + # trigger as the CTAs exit still gives it the launch overlap. + early_q = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + hint_blocks = cute.ceil_div(seq_lens[batch_id], BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, BLOCK_K * head_dim * (dtype.width // 8) + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # seq_lens must be read after the wait. It is a graph-stable buffer that + # on_update_kv_lens patches every step, and num_blocks bounds an unmasked + # block_table load and an unmasked score store, so a stale length is an + # out-of-bounds write rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score_t2sc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score_t2sc.py new file mode 100644 index 000000000000..9c0757ed166a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score_t2sc.py @@ -0,0 +1,682 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule: T2SC, this target's copy of +the sm_107 TP2 / TP8 targets' scorer (``modeling.INDEX_SCORE_T2SC``). The sm_103-derived ``kernels/m3_index_score`` +stays the kill path. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with three schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. +* ``prestage_old`` (with ``early_issue``, ``modeling.INDEX_SCORE_T2SC_PRESTAGE_OLD``): when the CTA's first block holds + only tokens older than this step (block < (seq_lens hint - decode query length) / 128), its K tile is TMA-loaded into + its stage before the wait; Q (and a first block that is not all old) is still loaded after it. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer_t2sc``) through +``IndexScoreRunner`` and called through this target's ``m3_index_score_t2sc`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import ( + TRTLLM_ENABLE_PDL, + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + prestage_old: bool = False, + early_trigger: bool = False, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K + # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's + # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. + self.prestage_old = prestage_old and early_issue + # early_trigger (on for this target): release the dependent selector right after this CTA's + # setup, so its launch and residency overlap the scoring; it still waits for this whole grid. + self.early_trigger = early_trigger + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # Without early_trigger the selector launches at this grid's implicit trigger (the CTAs' exit). With it, the + # selector is already resident in its own grid-dependency wait, which still returns only once every block + # this grid writes is complete. + early_q = False + early_k = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + seqlen_hint = seq_lens[batch_id] + hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + if cutlass.const_expr(self.prestage_old): + # Stage the first K tile now when every token of the block predates this step. + if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: + early_k = True + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + if not early_k: + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # The length that bounds this grid's work is read after the wait (the + # pre-wait reads above are prefetch / early-issue hints only, clamped to + # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens + # patches every step, and num_blocks bounds an unmasked block_table load + # and an unmasked score store, so a stale length is an out-of-bounds write + # rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + prestage_old=early, + early_trigger=True, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_lowm_splitk.py new file mode 100644 index 000000000000..bf20343f7886 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_lowm_splitk.py @@ -0,0 +1,1249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +This target's copy of FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the Eagle3 +drafter's low-M linears, ``kernels/m3_draft_lowm``): the same kernel and the same split order at every tactic, +bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's final +cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, the +SwiGLU prologue that follow. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 # capacity: fixed: a kernel tile constant + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 # capacity: fixed: an alignment + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 # capacity: fixed: pipeline stages +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 # capacity: fixed: pipeline stages + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= _MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES + base = SplitKTactic(mma_m, mma_n, split_k, min_stages) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(min_stages, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 # capacity: fixed: a split index, not a size + + +# The drafter-boundary push (``ar_push``, modeling.DRAFT_AR_PUSH): the exchange words of +# kernels/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.s32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + llvm.inline_asm( + None, + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + ar_push: bool = False, + ) -> None: + self.ar_push = ar_push + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64 = cutlass.Int64(0), + ar_flags: cutlass.Int64 = cutlass.Int64(0), + ar_rank: Int32 = Int32(0), + ar_world: Int32 = Int32(1), + ar_maxt: Int32 = Int32(8), + ar_tokens: Int32 = Int32(0), + ar_n: Int32 = Int32(0), + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel( + a, b, c, bias, ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_tokens, ar_n + ).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + ar_ws: cutlass.Int64, # ar_push: the boundary all-reduce's multicast exchange words (int32) + ar_flags: cutlass.Int64, # ar_push: its per-(token, cluster CTA) call flags (parity = flag & 1) + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, # the bucket's tokens + ar_tokens: Int32, # rows pushed (public M) + ar_n: Int32, # public N (the hidden size): the exchange row width + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + # ar_push: the owner's BF16 tile (kernel-M = output column, kernel-N = token), M-major. + sD = cute_ext.allocate( + mC.element_type, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=16, + ) + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sD, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl and not self.ar_push): + # ar_push: the dependent is a prepushed all-reduce, which skips its grid wait: it is released only after + # this grid's own grid wait and its pushes, in the owner's _store_tile (the prepushed >= 2 contract). + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sD: cute.Tensor, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + self._store_tile(thr_t2r, rAcc, rD, gD_epi, c_dtype) + if cutlass.const_expr(self.ar_push): + self._push_tile( + thr_t2r, + rD, + sD, + epi_tid, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + @cute.experimental.jit + def _store_tile(self, thr_t2r, rAcc, rD, gD_epi, c_dtype: cutlass.Constexpr): + """The owner rank's reduced FP32 tile: one cast to ``c_dtype`` and the store to global memory.""" + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + @cute.experimental.jit + def _push_tile( + self, + thr_t2r, + rD, + sD, + epi_tid: cutlass.Int32, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output + would store them (kernels/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" + cute_ext.partition_and_copy(thr_t2r, rD, sD) + cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + pairs = self.cta_m // 2 + col0 = bidx * self.cta_m + half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) + for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): + word = epi_tid + w0 + t = word // pairs + p = word % pairs + if t < ar_tokens: + n = col0 + 2 * p + par = _ld_volatile_s32( + ar_flags + cutlass.Int64(4) * cutlass.Int64(t * 8 + n // (ar_n // 8)) + ) & Int32(1) + off = ( + cutlass.Int64(par) * half + + cutlass.Int64(t * ar_world + ar_rank) * cutlass.Int64(ar_n // 2) + + cutlass.Int64(n // 2) + ) + _st_global_b32( + ar_ws + off * cutlass.Int64(4), + _bf16x2_word( + sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + ), + ) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +import torch as _torch # noqa: E402 (as in FlashInfer) + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ) + + +@cute.experimental.jit +def _bmm_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ar_tokens: Int32, + ar_n: Int32, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any(tensor.dtype != a.dtype for tensor in tensors): + raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a, + b, + bias, + out, + pdl: bool, + tactic: SplitKTactic, +): + """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +@functools.cache +def _get_compiled_push( + dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] +): + kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, ar_push=True) + compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + z64, z32 = cutlass.Int64(0), Int32(0) + return cute_ext.compile( + _bmm_push, + kernel, + *compile_tensors[:3], + stream, + z64, + z64, + z32, + Int32(1), + Int32(8), + z32, + Int32(6144), + ) + + +def compile_push(a, b, out, pdl: bool, tactic) -> None: + """Compile (only) the push variant of ``run_splitk_dense_push`` for these operands' shapes and layouts.""" + cute_tensors = _to_cute_swap(a, b, out, None) + _get_compiled_push(a.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) + + +def run_splitk_dense_push(a, b, out, pdl: bool, tactic: SplitKTactic, ar: tuple): + """``run_splitk_dense`` without bias that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange; + ``ar`` = (multicast words address, flags address, rank, world, bucket tokens) of the 8-token bucket.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_push(a.dtype, tactic, pdl, cute_tensors[4]) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + ws, flags, rank, world, maxt = ar + compiled( + *cute_tensors[:3], + stream, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + Int32(out.shape[0]), + Int32(out.shape[1]), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_mc_allreduce.py new file mode 100644 index 000000000000..142bcb9a9133 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_mc_allreduce.py @@ -0,0 +1,422 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-latency TP all-reduce + residual + RMSNorm for decode (M <= 8 tokens) over a multicast mapping. + + partial = bf16(x + x2) (optional second input, e.g. shared + routed experts) + updated = bf16(residual + sum_r partial_r) (ranks summed in FP32, rank order) + normed = bf16(updated * rsqrt(mean(updated^2) + eps) * w) + +One cluster of CLUSTER CTAs per token; CTA c owns columns [c * H / CLUSTER, ...), one 16-byte chunk (8 BF16) +per thread. Each CTA pushes its slice of this rank's partial through the multicast mapping (one store lands in +every rank's buffer), releases dependents, polls its own buffer until all ranks' slices arrived, sums them, adds +the residual and reduces the RMS statistic across the cluster through DSMEM. + +Exchange protocol as in the K3 sandwich: int32 words, 0x80000000 = empty (pushed -0.0 halves become +0.0), each +CTA keeps a call counter whose parity selects the buffer half, readers re-arm the words they read. All ranks of +the TP group must issue the same calls in the same stream order. + +The workspace holds two exchange sets. With prepushed 3 two producers pushed this rank's x (set A) and x2 (set B), +e.g. the routed experts' finalize and the shared expert's down projection, and the kernel forms bf16(x + x2) per +rank from the two slots, exactly as the unpushed path does. + +Producers that push (the o_proj / shared-expert GEMV ``ar`` tuple, the routed-expert engine's ``arp`` descriptor) are +pointed at a bucket's exchange by its ``mc`` address, ``flags_for(8)`` counters, ``rank``, ``world_size`` and +``max_tokens``; ``get_workspace`` and ``_bucket`` select it. + +The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp4_mc_allreduce_norm``, reached through the ``ops/comm`` entries +``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``; this module owns the workspaces, the bucket and +epilogue-option choice and the MXFP8 publication. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op +from ..ops.comm.m3_mc_finalize_allreduce_norm import ( + m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, +) + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +BUCKETS = (8, 16, 48) # token capacity of each workspace (grid = bucket * cluster CTAs) +EMPTY_WORD = -2147483648 + +# Attention / MoE boundary epilogue options (each bit-identical): 1 residual + norm weight loaded at entry (before +# the poll); 2 cluster RMS exchange by st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated +# right after the poll (no trailing __syncthreads); 8 max-smem carveout so the next kernel's CTAs co-reside early. +# The 8-token bucket, by token count 1..8 (GB300 TP4 tuning). +AR_OPTS_BY_M = (7, 7, 15, 15, 15, 15, 7, 7) +# The 16- and 48-token buckets. +AR_OPTS_BY_BUCKET = {16: 15, 48: 15} +# The MoE-finalize all-reduce, per bucket. +AR_FIN_OPTS_BY_BUCKET = {8: 7, 16: 15, 48: 15} + + +# Option 32: the 8-CTA cluster kernel in the stock TP4 oneshot reduction order instead of the row kernel (not finalize). +_OPT_ONESHOT_CLUSTER = 32 + +# Option 64: the row kernel's max-smem-carveout instance (as option 8 is for the cluster kernel). +_OPT_ROW_CARVEOUT = 64 + +# Option 128: the drafter boundary (option 16) on the 8-CTA cluster kernel instead of the world-4 row kernel (routing +# only; mcAllreduceNormTp4.cu invoke). Exact-class to the row kernel's FlashInfer order: the RMS statistic is the +# cluster's sum (modeling.DRAFT_MC_AR_CLUSTER). +_OPT_DRAFTER_CLUSTER = 128 + +# Option 2 is also the safe ordering of the cluster RMS exchange: its stores into peer CTAs' shared memory follow a +# cluster barrier arrive / wait pair. Without it they run before any cluster barrier, when a peer need not have +# started, so every call keeps it. +_OPT_CLUSTER_BARRIER = 2 + + +def _ar_opts(m: int, bucket: int, finalize: bool = False, oneshot_cluster: bool = False) -> int: + """Epilogue options for an m-row call on the given bucket; every option is bit-identical. A bucket missing from + the tables gets 0, and every result keeps _OPT_CLUSTER_BARRIER. ``oneshot_cluster`` adds _OPT_ONESHOT_CLUSTER.""" + if oneshot_cluster and not finalize: + return _ar_opts(m, bucket) | _OPT_ONESHOT_CLUSTER + if finalize: + opts = AR_FIN_OPTS_BY_BUCKET.get(bucket, 0) + elif bucket == 8: + opts = AR_OPTS_BY_M[m - 1] + else: + opts = AR_OPTS_BY_BUCKET.get(bucket, 0) + return opts | _OPT_CLUSTER_BARRIER + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags + + +class McBuckets: + """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" + + def __init__(self, workspaces): + self.workspaces = workspaces + self.hidden = workspaces[0].hidden + + def for_tokens(self, m: int): + for ws in self.workspaces: + if m <= ws.max_tokens: + return ws + return None + + +def _bucket(workspace, m): + return workspace.for_tokens(m) if isinstance(workspace, McBuckets) else workspace + + +def decline_reason(x, residual, workspace, cluster=8, x2=None) -> str: + """Why ``m3_mc_allreduce_norm`` returned None for these operands (for the dispatch's miss message).""" + if workspace is None: + return "no multicast workspace: it is created collectively on the first eager decode step" + bucket = _bucket(workspace, x.shape[0]) + if bucket is None: + return f"no workspace bucket holds {x.shape[0]} tokens (buckets {BUCKETS})" + if bucket.flags.get(cluster) is None: + return f"no {cluster}-CTA call counters for the {bucket.max_tokens}-token bucket" + return ( + f"operands outside the kernel: x {tuple(x.shape)} {x.dtype} contiguous={x.is_contiguous()}, " + f"residual contiguous={residual.is_contiguous()}, hidden {bucket.hidden}" + + ("" if x2 is None else f", x2 {tuple(x2.shape)} contiguous={x2.is_contiguous()}") + ) + + +def m3_mc_allreduce_norm( + x, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + x2=None, + quant=False, + prepushed=0, + oneshot_cluster=False, + drafter_norm=False, + carveout=False, + capture=None, + drafter_cluster=False, +): + """Return (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture`` = (out [>= M, H] view, fc_norm weight + [H], eps): the boundary also writes the Eagle3 capture of the updated residual, bit for bit + m3_spec_tail.capture_norm (CAPTURE_NORM_IN_BOUNDARY; not with ``drafter_norm``). ``drafter_cluster``: with + ``drafter_norm``, the 8-CTA cluster kernel's drafter form (option 128; its carveout is option 8) instead of the row + kernel's. + + A None makes the caller take another collective, so it must be the same on every TP rank. It is: it depends only + on the shapes, dtypes and layouts of the operands (the same code builds them on every rank) and on whether the + workspace's flags exist, which ``get_workspace`` creates collectively, before any capture, for cluster 8.""" + workspace = _bucket(workspace, x.shape[0]) + if workspace is None: + return None + if ( + x.dim() != 2 + or x.shape[0] > MAX_TOKENS + or x.dtype != torch.bfloat16 + or x.shape[1] != workspace.hidden + or not x.is_contiguous() + or not residual.is_contiguous() + or (x2 is not None and (not x2.is_contiguous() or x2.shape != x.shape)) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + # The Eagle3 drafter boundary (opts bit 16, the row kernel's FlashInfer fused_add_rmsnorm order): 8-token bucket, + # x only, no MXFP8 copy, never the cluster kernel. + if drafter_norm and ( + workspace.max_tokens != 8 or x2 is not None or quant or prepushed not in (0, 2) + ): + return None + if drafter_norm and drafter_cluster: + opts = ( + (16 | _ar_opts(x.shape[0], workspace.max_tokens)) & ~8 + | (8 if carveout else 0) + | _OPT_DRAFTER_CLUSTER + ) + elif drafter_norm: + opts = ( + 16 | _ar_opts(x.shape[0], workspace.max_tokens) | (_OPT_ROW_CARVEOUT if carveout else 0) + ) + else: + opts = _ar_opts(x.shape[0], workspace.max_tokens, oneshot_cluster=oneshot_cluster) + if carveout: + # the max-smem carveout on whichever kernel runs: option 8 on the cluster kernel, 64 on the row kernel + opts |= 8 if oneshot_cluster else _OPT_ROW_CARVEOUT + outs = _allreduce_norm_op( + x, + x2, + residual, + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and x.shape[1] % 128 == 0), + int(prepushed), + opts, + *_capture_args(capture, drafter_norm), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def _capture_args(capture, drafter_norm: bool) -> tuple: + """(cap_out, cap_w, cap_eps) for the op, from ``capture`` = (out, fc_norm weight, eps) or None.""" + if capture is None: + return (None, None, 0.0) + assert not drafter_norm, "the capture epilogue is not built for the drafter boundary" + out, weight, eps = capture + return (out, weight, float(eps)) + + +def m3_mc_finalize_allreduce_norm( + fc2, + perm, + scales, + shared, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + quant=False, +): + """MoE finalize (sum_k scale * fc2[perm] + shared) + all-reduce + residual + RMSNorm in one kernel. + Returns (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. (It runs the shared GB300 op, which has no capture + epilogue: a capture layer finalized here keeps the capture_norm kernel.)""" + m = residual.shape[0] + workspace = _bucket(workspace, m) + if workspace is None: + return None + if ( + residual.dim() != 2 + or m > MAX_TOKENS + or residual.shape[1] != workspace.hidden + or fc2.dtype != torch.bfloat16 + or not fc2.is_contiguous() + or fc2.dim() != 2 + or perm.dtype != torch.int32 + or perm.dim() != 2 + or perm.shape[0] != m + or not perm.is_contiguous() + or scales.dtype not in (torch.float32, torch.bfloat16) + or not scales.is_contiguous() + or (shared is not None and (shared.shape != residual.shape or not shared.is_contiguous())) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _finalize_allreduce_norm_op( + fc2, + perm, + scales, + shared, + residual.contiguous(), + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and residual.shape[1] % 128 == 0), + _ar_opts(m, workspace.max_tokens, finalize=True), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +_WORKSPACE: dict = {} +_AGREE_COMMS: dict = {} + + +def _tp_comm(mapping): + """A communicator over mapping's TP group, created once (collective over that group only).""" + key = tuple(mapping.tp_group) + comm = _AGREE_COMMS.get(key) + if comm is None: + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + comm = mapping.tp_group_pg + else: + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + _AGREE_COMMS[key] = comm + return comm + + +def tp_agree(mapping, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of mapping's TP group (one blocking host collective). + + For a fast path whose availability is decided from rank-local state (a kernel build that can fail on one rank): + every rank must take the same path, so that they issue the same collectives with the same exchange protocol and + the same kernels. Call it once per decision, at build time, at the same point on every rank, outside CUDA-graph + capture; never per step. With one rank, or attention DP (whose ranks run different batches), it returns ``ok``. + """ + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return ok + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("tp_agree must run outside CUDA-graph capture") + from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded + + return _mnnvl_workspace_all_succeeded(_tp_comm(mapping), bool(ok)) + + +def tp_agree_built(mapping, built, what: str): + """``built`` when it is not None on every TP rank, else None on every rank (see ``tp_agree``).""" + if tp_agree(mapping, built is not None) or built is None: + return built + from tensorrt_llm.logger import logger + + logger.warning(f"{what}: built here but not on another TP rank; every rank falls back") + return None + + +def get_workspace(mapping, hidden: int): + """Process-wide workspace (created collectively on first eager use); None while capturing before creation. + + The first eager call for a hidden size (the warmup forward, before any CUDA-graph capture) creates one + McAllReduceWorkspace per bucket in BUCKETS on every TP rank (MPI collective over the TP group) plus each + bucket's 8-CTA-cluster call counters; later calls, captured or not, return the same McBuckets. The workspaces + live until process exit. + """ + ws = _WORKSPACE.get(hidden) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + return None + wss = [McAllReduceWorkspace.create(mapping, hidden, b) for b in BUCKETS] + for w in wss: + w.flags_for(8) + ws = McBuckets(wss) + _WORKSPACE[hidden] = ws + return ws diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/__init__.py new file mode 100644 index 000000000000..f818a0e6c9f1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/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 sm_107 TP4 target for low-batch decode as a CuTe DSL kernel: ``m3_moe`` for 1..8 tokens +(op.py) at the TP4 routed and shared intermediate of 768 per rank.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe16_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe16_kernel.py new file mode 100644 index 000000000000..ff609e78ef29 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe16_kernel.py @@ -0,0 +1,3871 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 routed + shared experts for 9..16 decode tokens on one TP rank: ``m3_moe16`` (v2). + +Same math as the 1..8-token engine (m3_moe_kernel.py), operation for operation for the routed part (bitwise equal to +TRTLLM-Gen's fused MoE on the fp32 router-GEMV logits); a different schedule for 9..16 rows: + +- Front: the hidden states land in smem by bulk copies in K chunks (the idle ring); all 16 warps run the router + GEMV (m3_router_gemv's per-thread order) and the NVFP4 quantization of B1 from smem; the logits cross CTAs as + tagged words; top-k one warp per token. +- Ring: FC1 / FC2 weight stages by TMA; their scale words go by 4-byte cp.async straight into the instruction atom + layout (SFA_CP; no raw scale box, no splicer warp), completing on the stage's full barrier. +- FC1 -> FC2 hand-off (B2IMG): the FC1 epilogue writes each unit's FP4 values and scale byte into a per-group global + image laid out as the FC2 B tile + its scale atoms; one counter per group (release add per finished unit, every + counter gains exactly 192 per call); warp 5 loads a completed image with two bulk copies. +- Shared expert on NSH more CTAs (MXFP8 tcgen05), as the 1..8 engine. +Configuration via M3_CONFIG (m_max, num_ctas, stages, nb2_cap, lag, shared, sfa_cp, b2buf, ...). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +H = 6144 # capacity: fixed: the checkpoint's hidden size +I_TP = 768 +TWO_I = 2 * I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 12) +assert 9 <= M_MAX <= 16, M_MAX +NCTA = _cfg("num_ctas", 128) +FUSED_ROUTER = True +NROUTE = E_LOCAL +NGEMV = (NROUTE + NCTA - 1) // NCTA +assert NGEMV == 1, "m3_moe16: one router expert per CTA (num_ctas >= 128)" +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +LAG = _cfg("lag", 3) +G_CAP = TOP_K * M_MAX +NG_CAP = M_MAX +THREADS = 512 # capacity: fixed: the CTA size +NWARPS = THREADS // 32 +GEMV_THREADS = 256 +SHARED = bool(_cfg("shared", 0)) +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +SFA_CP = bool( + _cfg("sfa_cp", 1) +) # A scales by 4-byte cp.async into the atom layout (else raw box + splicer warp) +B2IMG = True +PERM = bool(_cfg("perm", 1)) +PERM_MIN_NG = max( + 3, _cfg("perm_min_ng", 5) +) # (measured: helps at random routing G >= 23, hurts at spec G 13-14) +# a partial last slot group becomes group 1 (the plan must match: op16) +DBG_DELAY = bool( + _cfg("dbg_delay", 0) +) # (stress builds) pseudo-random producer / epilogue / release stalls +DBG_NOSTORE = bool(_cfg("dbg_nostore", 0)) # (timing experiment only: wrong results) +EPI_FENCE = bool( + _cfg("epi_fence", 0) +) # (rel_warp 0) writer lanes' fence.acq_rel.gpu before the release add +EPI_PFENCE = bool( + _cfg("epi_pfence", 1) +) # producer-side fence.proxy.async.global before the release add +REL_WARP = bool( + _cfg("rel_warp", 1) +) # warp 8 does the per-group release adds (off the epilogue's path) +NSH = 24 if SHARED else 0 +SH_I = 768 +SH_KB = H // 32 # capacity: fixed: a kernel tile constant +SH_DKB = SH_I // 32 # capacity: fixed: a kernel tile constant +SH_SPLIT = 2 +SH_GT = SH_KB // 4 // SH_SPLIT # capacity: fixed: a kernel tile constant +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * (SH_I // 128) # capacity: fixed: pipeline stages +SHX_BAR = 6 +SH_PRETILE = False +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = _cfg("n1", 48 if _cfg("m_max", 12) <= 12 else 64) # B rows QS q + t +N1_LD = 1 << (N1 - 1).bit_length() # tcgen05.ld .num (a power of two) covering N1 columns +QS = N1 // 4 +assert QS >= M_MAX +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 +K1_TILES = H // 256 +S1 = K1_TILES // C1 # capacity: fixed: pipeline stages +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 +KT2 = I_TP // 256 +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = _cfg("n2", (4 * M_MAX + 15) // 16 * 16) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +IMG_B2 = KT2 * NUM_BYTES_B2 # one group's B2 tile image (global and smem) +IMG_SF = KT2 * NUM_KB * 512 # its scale atoms +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) # every group counter gains 192 per call +B1IMG_OFF = None # set below (after N1): the activation image after the NG_CAP group images +TK_WARPS = [0, 1, 2, 3, 6, 8, 9, 10, 11, 12, 13, 14, 5, 7, 15, 4][ + : _cfg("m_max", 12) +] # top-k warp of token t (4 last: early0) +B1_WARP = 15 if _cfg("m_max", 12) <= 14 else -1 # (-1: warp 0 after the routing) +NB2_CAP = _cfg("nb2_cap", 2) +B1IMG_OFF = M_MAX * IMG_BYTES +B1IMG_BYTES = (S1 * N1 * 128 + M_MAX * (H // 16) + 15) // 16 * 16 +IMG_TOTAL = B1IMG_OFF + B1IMG_BYTES +UNITS_ROW = TWO_I // R1 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +TOPK_BAR = 4 +ALL_BAR = 5 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 16) +LG_REPL = _cfg("lg_repl", 64) +LG_STRIDE = M_MAX * E_LOCAL + 32 +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # B2 image loader; slots + plan # capacity: fixed: a warp role index +S_WARP = 6 # (SFA_CP = 0) scale splicer # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +EARLY0 = bool(_cfg("early0", 1)) +EVICT_FIRST = 0x12F0000000000000 +EVICT_LAST = 0x14F0000000000000 +_LOG2E = 1.4426950408889634 +NB2BUF = _cfg("b2buf", 2) +SPB = 1 +# x staging: K chunks of XCK elements, NXB smem buffers (rotating if fewer than the chunks) +XCK = _cfg("xck", 2048) +NXC = H // XCK +XCB = M_MAX * XCK * 2 # bytes of one chunk (all tokens) +STAGE_A = NUM_BYTES_A +STAGE_SF = NUM_BYTES_SFA if SFA_CP else NUM_BYTES_SFA_RAW + NUM_BYTES_SFA +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + NB2BUF * IMG_B2 + + IMG_SF + + NB2_CAP * M_MAX * TOP_K * R2 * 2 + + SPB * M_MAX * C1 * R1 * 4 + + (G_CAP + 1) * (NG_CAP + 1) * 4 + + 6 * 1024 +) +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // (STAGE_A + STAGE_SF) +ACC1_COL = 0 +ACC2_COL = 2 * N1 +ACC2_STRIDE = 1 << (N2 - 1).bit_length() +SFB1_COL = ACC2_COL + 2 * ACC2_STRIDE +SFB1_WSTEP = 1 +SFB2_COL = ( + SFB1_COL + NUM_SFB1_ATOMS * 4 +) # one TMEM slot (loaded after the previous group's FC2 MMAs) +SFA1_COL = SFB2_COL + KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# smem: [ring A | ring SF | B2 bufs | B2 SF staging | B1]; x chunks alias [0, _OFF_B1) during the front +_OFF_SFA = STAGES * NUM_BYTES_A +_OFF_B2 = (_OFF_SFA + STAGES * STAGE_SF + 1023) // 1024 * 1024 +_OFF_SFB2 = _OFF_B2 + NB2BUF * IMG_B2 +_OFF_B1 = (_OFF_SFB2 + IMG_SF + 1023) // 1024 * 1024 +BIG = _OFF_B1 + S1 * NUM_BYTES_B1 +NXB = min(NXC, _OFF_B1 // XCB) +assert NXB >= 2, (NXB, _OFF_B1, XCB) +# shared CTA layout in BIG: one pass over all M tokens; B rows SHN q + t (gate_up, 2 K slices), t (down) +SHN = 16 # capacity: fixed: a kernel tile constant +SH_NPASS = 1 +SH_HINT = EVICT_FIRST +SH_PF_POST = bool( + _cfg("sh_pf_post", 1) +) # L2 prefetch of the shared weights after the grid dependency (else before) +_SB_A = 0 +_SH_FIX = ( + SH_GT * 2 * SHN * 128 + + 6 * SHN * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + 64 * SHN * 4 + + 2 * 32 * SHN * 4 + + 2048 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * 2 * SHN * 128 +_SB_XC = _SB_BH + 6 * SHN * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + 64 * SHN * 4 +_SB_END = _SB_GU + 2 * 32 * SHN * 4 +SH_ACC_GU = 0 +SH_ACC_DN = 2 * SHN +SH_SFA = SH_ACC_DN + 2 * SHN +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT # capacity: fixed: a kernel tile constant +SH_TMEM = 256 # capacity: fixed: TMEM columns +if SHARED: + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) +TD_STRIDE = 64 +PFD = _cfg("pfd", 0) # L2 prefetch distance in ring stages (0: off) +PF_WARP = 12 # capacity: fixed: a warp role index +TICKET = bool(_cfg("ticket", 0)) # CTA roles by entry order +# priority tokens: token 0 and PT1 (M / 2: the second request's first row at spec c2) get their logits published first +NPT = _cfg("npt", 0) +PT1 = M_MAX // 2 +PF_PRIO = bool( + _cfg("pf_prio", 0) +) # L2 prefetch of the priority tokens' experts (this CTA's 1/NCTA of each) +PF_W1 = TWO_I * (H // 2) // NCTA +PF_S1 = TWO_I * (H // 16) // NCTA +PF_W2 = H * (I_TP // 2) // NCTA +PF_S2 = H * (I_TP // 16) // NCTA +assert PF_W1 % 16 == 0 and PF_S1 % 16 == 0 and PF_W2 % 16 == 0 and PF_S2 % 16 == 0 + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + # B rows n >= 32 (N1 = 64) go to the atom's next 32-bit column (row n: lane n % 32, column n // 32) + for w in cutlass.range_constexpr((N1 + 31) // 32): + n = lane + cutlass.Int32(32 * w) + q = n // cutlass.Int32(QS) + tl = n - q * cutlass.Int32(QS) + sf_ok = (n < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr( + base_row + warp * cutlass.Int32(32), + base_col + SFB1_COL + a * 4 + w * SFB1_WSTEP, + ), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + gpos, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = gpos % cutlass.Int32(SH_NST) # ring position (stage g of this pass) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue( + lane, + sid, + cutlass.Int32(g), + s_big, + sh_full, + tma_gu, + tma_dn, + gusf, + dnsf, + cutlass.Int32(g), + ) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + elif cutlass.const_expr(not SH_PF_POST): + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(SH_PF_POST): + if warp == 0: # after the grid dependency: DRAM is idle during the routed front + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + for ps in cutlass.range_constexpr(SH_NPASS): + TT0 = 0 # one pass: B rows SHN q + t (gate_up), t (down) + MS = M_MAX + GB = ps * SH_STAGES # ring positions before this pass + PAR = ps & 1 # parity of the per-pass barriers (B ready, accumulators) + if warp == 0: + for g in cutlass.range(SH_NST if ps == 0 else 0, SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) # ring position over the passes + _wait( + sh_empty.subview(gg % cutlass.Int32(SH_NST)), + ((gg // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf, gg) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=2 * SHN, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SHN, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) + st = gg % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (gg // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32((2 * SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32((SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(SHN), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + if cutlass.const_expr(ps > 0): + _wait( + acc_dn, PAR ^ 1 + ) # the previous pass's MMAs are done: its B tiles and TMEM scales are free + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((MS * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(MS * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=(t + cutlass.Int32(TT0)) * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(SHN) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + + it * cutlass.Int32(2 * SHN * 128) + + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of + # k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(SHN - 1) + ok = (n < cutlass.Int32(2 * SHN)) & (t < cutlass.Int32(MS)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n // cutlass.Int32(SHN)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t + # at column 8 q + t + _wait(acc_gu, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=2 * SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(MS): + bigf.store( + cutlass.Float32(av[SHN + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(MS): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((32 * MS + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(32 * MS): + p = pq // cutlass.Int32(MS) + tt = pq - p * cutlass.Int32(MS) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(SHN) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(32) + p) * cutlass.Int32(SHN) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview( + (tt + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + + sid * cutlass.Int32(32) + + p + ) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((MS * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(MS * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = (t + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + kb * cutlass.Int32( + 32 + ) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + + kt * cutlass.Int32(SHN * 128) + + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (t & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + ( + ((c0 + cutlass.Int32(1)) ^ (t & cutlass.Int32(7))) + << cutlass.Int32(2) + ), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(MS) + for kt in cutlass.range_constexpr(6): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * SHN, + cutlass.Float32, + ), + num=SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(MS): + _st_relaxed_u32( + dnx.subview(cutlass.Int32((tt + TT0) * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +# ============================================================================= +# v2 helpers: global release / acquire, proxy fences. +# ============================================================================= +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _atom_add_u64(addr, val, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr), cutlass.Int64(val)], + "atom.relaxed.gpu.global.add.u64 $0, [$1], $2;", + "=l,l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.proxy.async.global;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_acq_rel_gpu(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.acq_rel.gpu;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _dbg_delay(salt): + """(DBG_DELAY stress builds) a pseudo-random ~0..4 us stall: protocol robustness under skew.""" + if cutlass.const_expr(DBG_DELAY): + r = cutlass.Int32(_gtime() >> cutlass.Int64(6)) * cutlass.Int32(1103515245) + salt + if ((r >> cutlass.Int32(12)) & cutlass.Int32(7)) == cutlass.Int32(0): + for _ in cutlass.range((r >> cutlass.Int32(16)) & cutlass.Int32(31), unroll=1): + _nanosleep(128) + + +@cute.jit +def _waitt(bar: cutlass.Array, parity, acc): + """_wait that (TIMING) adds the waited ns to acc.""" + t0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + if cutlass.const_expr(TIMING): + acc = acc + (_gtime() - t0) + return acc + + +@cute.jit +def _tstore(tdbg: cutlass.Array, bx, i: cutlass.Constexpr, v): + if cutlass.const_expr(TIMING): + tdbg.store(cutlass.Int64(v), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@cute.jit +def _issue_fc1( + lane, + stage, + e, + c, + j, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a1, + tma_s1, + s1w: cutlass.Array, +): + """FC1 stage j of unit c of expert e into ring stage `stage` (the caller waited its empty barrier): 4 TMA boxes + (rows 32 c .., k-tiles 4 j + q) and the stage's scale atoms (SFA_CP: word [r][q] of atom kb <- raw atom + (k-tile 4 j + q, kb) word [r][c % 4], by this warp's 32 lanes).""" + row = c * cutlass.Int32(R1) + kx = j * cutlass.Int32(C1 * KT_BYTES) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview(stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(q * NUM_BYTES_BOX1)), + tma_a1, + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32(SF1_KA) + + j * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview(stage * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1, + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + atom0 = (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32( + SF1_KA + ) + j * cutlass.Int32(16) + src0 = atom0 * cutlass.Int32(128) + lane * cutlass.Int32(4) + (c & cutlass.Int32(3)) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _issue_fc2( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a2, + tma_s2, + s2w: cutlass.Array, +): + """FC2 k-tile t2 of output block blk for the group's n_here slots: box jj = rows 32 blk .. of slot jj's w2 at + k-tile t2; scale word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) word [r][blk % 4].""" + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32(NUM_BYTES_BOX1 if SFA_CP else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2, + (cutlass.Int32(t2 * KT_BYTES), blk * cutlass.Int32(R2), e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2, + ( + cutlass.Int32(0), + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _quant_words(w, sc, rsc, rsix): + """NVFP4 of one 16-value block given as 8 bf16-pair words: (lo word, hi word, scale e4m3 bits).""" + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.jit +def _quant_block(xs32: cutlass.Array, widx, sc, rsc, rsix): + """NVFP4 of one 16-value block (8 bf16 words at widx): (lo word, hi word, scale e4m3 bits), fp4_quantize's ops.""" + w = [] + for h in cutlass.range_constexpr(2): + v4 = xs32.load(idx=widx + cutlass.Int32(4 * h), vector_size=4, alignment=16) + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_gu_desc: cutlass.GridConstant[cuda.TensorMap], + tma_dn_desc: cutlass.GridConstant[cuda.TensorMap], + bias: cutlass.Array, # f32 [128] + escal: cutlass.Array, # f32 [6, 128] + plan: cutlass.Array, # int32 [4 M + 1, NCTA, M + 1] + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + img32: cutlass.Array, # int32 [NG_CAP, IMG_BYTES / 4]: per-group FC2 B image (tile + scale atoms) + cnt: cutlass.Array, # int32 [NG_CAP * 8]: per-group unit counters (index 8 g) + s1w: cutlass.Array, # int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # int32 words of w2_weight_scale + w1b: cutlass.Array, # int8 bytes of w3_w1_weight (L2 prefetch) + w2b: cutlass.Array, # int8 bytes of w2_weight (L2 prefetch) + epochs: cutlass.Array, # int32 [NCTA + NSH] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8] + tdbg: cutlass.Array, # int64 [NCTA + NSH, 64] + xb32: cutlass.Array, # int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # f32 router weight [128, 6144] + sfsc: cutlass.Array, # f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # int64 [LG_REPL, LG_STRIDE] + guw: cutlass.Array, + gusf: cutlass.Array, + dnw: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if cutlass.const_expr(TICKET): + # role by entry order: the CTAs that become resident last (behind the predecessor's CTAs) take the shared + # expert roles (most slack), the first NCTA the routed ones (the router GEMV / logits exchange) + s_tk = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + if tidx == 0: + old = _atom_add_u64( + cnt.subview(cutlass.Int32(NG_CAP * 8 + 8)).data_ptr().toint(), cutlass.Int64(1) + ) + s_tk.store(cutlass.Int32(old % cutlass.Int64(NCTA + NSH)), idx=0) + prims.barrier_cta_sync(0) + bx = s_tk.load(idx=0) + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + # ring scale words: SFA_CP: the instruction atoms [STAGES][4 x 512 B]; else raw boxes [STAGES][8 KB] + spliced + sSF = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(STAGES * STAGE_SF // 4,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA + (0 if SFA_CP else STAGES * NUM_BYTES_SFA_RAW)), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=512, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(NB2BUF * IMG_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), shape=(IMG_SF // 4,), dtype=cutlass.Int32, alignment=1024 + ) + ys = cutlass.Array( + cutlass.BFloat16, + NB2_CAP * M_MAX * TOP_K * R2, + space=cutlass.AddressSpace.smem, + alignment=16, + ) # bf16 values + s_part = cutlass.Array( + cutlass.Float32, SPB * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array( + cutlass.Float32, 9 * M_MAX + 16, space=cutlass.AddressSpace.smem, alignment=16 + ) + # the activation's scale codes [M, 384] until they are in TMEM: aliases ys (first written by the FC2 epilogue) + s_xsf = cutlass.Array( + ys.data_ptr(0), shape=(M_MAX * (H // 64),), dtype=cutlass.Int32, alignment=16 + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + img_bar = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + sfb2_free = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_full = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + x_empty = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + route_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + ys_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + grp_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_load = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + topk_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tok0_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + px_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + s_prog = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array(cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_gn = cutlass.Array( + cutlass.Int32, NG_CAP, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots per group + s_emap = cutlass.Array(cutlass.Int32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array(cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32) + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + if tidx < NB2BUF: + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), 1) + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < NXB: + prims.mbarrier_init(x_full.subview(tidx), 1) + prims.mbarrier_init(x_empty.subview(tidx), NWARPS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + prims.mbarrier_init(sfb2_free.subview(0), 1) + prims.mbarrier_init(route_bar.subview(0), 32) + prims.mbarrier_init(ys_done.subview(0), EPI_THREADS) + prims.mbarrier_init(b1_load.subview(0), 1) + prims.mbarrier_init(topk_bar.subview(0), M_MAX * TOP_K) + prims.mbarrier_init(tok0_bar.subview(0), TOP_K) + prims.mbarrier_init(px_full.subview(0), 1) + if tidx < 4: + prims.mbarrier_init(grp_done.subview(tidx), EPI_THREADS) + if tidx == 0: + s_prog.store(cutlass.Int32(0), idx=0) + if warp == X_WARP: # slot tables: empty (the slot warp fills them after the top-k) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + for z in cutlass.range_constexpr(E_LOCAL // 32): + s_emap.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + if cutlass.const_expr(not SFA_CP): + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + # router weights (expert bx) in registers: thread tg = tid % 256 owns elements 4 tg + 1024 j + tg = tidx % cutlass.Int32(GEMV_THREADS) + half = tidx // cutlass.Int32(GEMV_THREADS) + rwv = [ + rw.load( + idx=bx * cutlass.Int32(H) + tg * cutlass.Int32(4) + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + # ---- x chunks into the idle ring (warp 6): chunk c = K [XCK c, XCK (c + 1)) of every token + if warp == S_WARP: + if cutlass.const_expr(NPT > 0): + # the priority tokens' rows first (their logits are published ahead of the others) + if lane == 0: + prims.mbarrier_arrive_expect_tx(px_full.subview(0), NPT * H * 2) + cute.arch.sync_warp() + if lane < NPT: + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(lane * cutlass.Int32(H * 2)), + xb32.subview( + cutlass.select_( + lane == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1) + ) + * cutlass.Int32(H // 2) + ), + px_full.subview(0), + H * 2, + ) + for c in cutlass.range_constexpr(NXB): + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(c), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(c * XCB) + lane * cutlass.Int32(XCK * 2)), + xb32.subview(lane * cutlass.Int32(H // 2) + cutlass.Int32(c * XCK // 2)), + x_full.subview(c), + XCK * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + img8f = cutlass.Array( + img32.data_ptr(0), shape=(IMG_TOTAL,), dtype=cutlass.Int8, alignment=16 + ) + if cutlass.const_expr(NPT > 0): + # ---- priority tokens (0 and PT1): their logits first (m3_router_gemv's per-thread / warp order) + _wait(px_full.subview(0), 0) + ph = cutlass.select_(half < cutlass.Int32(NPT), half, cutlass.Int32(0)) + pacc = cutlass.Float32(0.0) + px32 = cutlass.Array( + sB1.data_ptr(0), shape=(NPT * H // 2,), dtype=cutlass.Int32, alignment=16 + ) + for j in cutlass.range_constexpr(H // 1024): + pxw = px32.load( + idx=ph * cutlass.Int32(H // 2) + cutlass.Int32(512 * j) + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + px0 = cutlass.Int32(pxw[0]) + px1 = cutlass.Int32(pxw[1]) + pacc = _fma_rn( + (px0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][0]), + pacc, + ) + pacc = _fma_rn( + (px0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][1]), + pacc, + ) + pacc = _fma_rn( + (px1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][2]), + pacc, + ) + pacc = _fma_rn( + (px1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][3]), + pacc, + ) + for off in (16, 8, 4, 2, 1): + pacc = _add_rn(pacc, cute.arch.shuffle_sync_bfly(pacc, off)) + if lane == 0: + s_red.store(pacc, idx=cutlass.Int32(9 * M_MAX) + warp) # [half][8 warps] + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if tidx < cutlass.Int32(NPT * LG_REPL): + hh = tidx // cutlass.Int32(LG_REPL) + rr = tidx - hh * cutlass.Int32(LG_REPL) + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn( + sacc, s_red.load(idx=cutlass.Int32(9 * M_MAX + w) + hh * cutlass.Int32(8)) + ) + tp = cutlass.select_(hh == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1)) + _st_relaxed_u64( + lgbuf.subview(rr * cutlass.Int32(LG_STRIDE) + tp * cutlass.Int32(E_LOCAL) + bx) + .data_ptr() + .toint(), + ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)), + ) + if tidx == 0: + _stamp(tdbg, bx, 43) + # ---- this CTA's share of the NVFP4 activation (QB blocks, from global) for the global B1 image: loads first + QB = (M_MAX * (H // 16) + NCTA - 1) // NCTA + qb = bx * cutlass.Int32(QB) + tidx + q_ok = (tidx < cutlass.Int32(QB)) & (qb < cutlass.Int32(M_MAX * (H // 16))) + qb = cutlass.select_(q_ok, qb, cutlass.Int32(0)) + qtok = qb // cutlass.Int32(H // 16) + qkb = qb - qtok * cutlass.Int32(H // 16) + xq_v = [ + xb32.load( + idx=qtok * cutlass.Int32(H // 2) + qkb * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + for h in range(2) + ] + # ---- router GEMV from the staged chunks, all 512 threads (rolled: the front runs once per call) + MH = (M_MAX + 1) // 2 # tokens per GEMV half + acc_r = cute.make_rmem_tensor((MH,), cutlass.Float32) + for i in cutlass.range_constexpr(MH): + acc_r[i] = cutlass.Float32(0.0) + big32 = cutlass.Array( + s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16 + ) + JPC = XCK // 1024 # router j-steps per chunk + for xc in cutlass.range(NXC, unroll=1): + buf = xc % cutlass.Int32(NXB) + _wait(x_full.subview(buf), (xc // cutlass.Int32(NXB)) & cutlass.Int32(1)) + if (tidx == 0) & (xc == cutlass.Int32(0)): + _stamp(tdbg, bx, 41) + xw0 = buf * cutlass.Int32(XCB // 4) # chunk word base in s_big + for jj in cutlass.range_constexpr(JPC): + jg = xc * cutlass.Int32(JPC) + cutlass.Int32(jj) + wsel = [cutlass.Float32(rwv[0][z]) for z in range(4)] + for jx in cutlass.range_constexpr(1, H // 1024): + for z in cutlass.range_constexpr(4): + wsel[z] = cutlass.select_( + jg == cutlass.Int32(jx), cutlass.Float32(rwv[jx][z]), wsel[z] + ) + for i in cutlass.range_constexpr(MH): + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + xw = big32.load( + idx=xw0 + + t * cutlass.Int32(XCK // 2) + + cutlass.Int32(512 * jj) + + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + acc_g = acc_r[i] + acc_g = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[0], acc_g + ) + acc_g = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[1], acc_g + ) + acc_g = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[2], acc_g + ) + acc_g = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[3], acc_g + ) + acc_r[i] = acc_g + if cutlass.const_expr(NXB < NXC): + if xc + cutlass.Int32(NXB) < cutlass.Int32(NXC): + # every thread's reads of this chunk are done (CTA barrier, then a generic -> async proxy fence by + # the refilling warp: the bulk copy overwrites what the generic proxy read); racecheck-clean + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == S_WARP: + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(buf), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + buf * cutlass.Int32(XCB) + lane * cutlass.Int32(XCK * 2) + ), + xb32.subview( + lane * cutlass.Int32(H // 2) + + (xc + cutlass.Int32(NXB)) * cutlass.Int32(XCK // 2) + ), + x_full.subview(buf), + XCK * 2, + ) + if tidx == 0: + _stamp(tdbg, bx, 42) + for i in cutlass.range_constexpr(MH): + v = acc_r[i] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + if (lane == 0) & (t < cutlass.Int32(M_MAX)): + s_red.store(v, idx=(warp % cutlass.Int32(8)) * cutlass.Int32(M_MAX) + t) + # this CTA's quantization blocks into the global B1 image (B1 tile bytes, then the [M, 384] scale codes) + if q_ok: + qw = [] + for h in cutlass.range_constexpr(2): + for z in cutlass.range_constexpr(4): + qw.append(cutlass.Int32(xq_v[h][z])) + qlo, qhi, qsf = _quant_words(qw, sc, rsc, rsix) + img32.store( + (qlo, qhi), idx=cutlass.Int32(B1IMG_OFF // 4) + _b1_addr(qtok, qkb), alignment=8 + ) + img8f.store(cutlass.Int8(qsf), idx=cutlass.Int32(B1IMG_OFF + S1 * NUM_BYTES_B1) + qb) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == cutlass.Int32(0): + _dbg_delay(bx * cutlass.Int32(3)) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX + t))) + lw = ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)) + _st_relaxed_u64( + lgbuf.subview(tidx * cutlass.Int32(LG_STRIDE) + cutlass.Int32(t * E_LOCAL) + bx) + .data_ptr() + .toint(), + lw, + ) + if tidx == 0: + _stamp(tdbg, bx, 21) + if tidx == cutlass.Int32( + THREADS - 32 + ): # B1 share released after the logits (its gpu fence stalls warp 15 only) + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() # once, after the join (ALL_BAR) of the writer threads + _red_release_add( + cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint(), cutlass.Int32(1) + ) + # ---- the B1 image (every CTA's share) into smem (warp 13; the top-k warps are 0 .. M - 1) + if cutlass.const_expr(B1_WARP >= 0): + if warp == cutlass.Int32(B1_WARP): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + + # ---- top-k: token t on warp t (keys sigmoid + bias, ties to the lower id), TRTLLM-Gen MiniMax2 ops + tk_t = cutlass.Int32(-1) # this warp's top-k token (TK_WARPS[t] = warp) + for i_tk in cutlass.range_constexpr(M_MAX): + tk_t = cutlass.select_(warp == cutlass.Int32(TK_WARPS[i_tk]), cutlass.Int32(i_tk), tk_t) + if tk_t >= cutlass.Int32(0): + t = tk_t + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 22) + lgv = [cutlass.Int32(res[z]).bitcast(cutlass.Float32) for z in range(4)] + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hk = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hk[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + if lane < TOP_K: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 20) + if cutlass.const_expr(PF_PRIO): + if (t == cutlass.Int32(0)) | (t == cutlass.Int32(PT1)): + if lane < TOP_K: + # this CTA's 1/NCTA of each of the token's experts into L2 (DRAM is idle until the routing is + # complete; every byte is streamed later anyway) + _l2_prefetch( + w1b.subview( + e_l * cutlass.Int32(TWO_I * H // 2) + bx * cutlass.Int32(PF_W1) + ) + .data_ptr() + .toint(), + PF_W1, + ) + _l2_prefetch( + s1w.subview( + (e_l * cutlass.Int32(TWO_I * H // 16) + bx * cutlass.Int32(PF_S1)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S1, + ) + _l2_prefetch( + w2b.subview( + e_l * cutlass.Int32(H * I_TP // 2) + bx * cutlass.Int32(PF_W2) + ) + .data_ptr() + .toint(), + PF_W2, + ) + _l2_prefetch( + s2w.subview( + (e_l * cutlass.Int32(H * I_TP // 16) + bx * cutlass.Int32(PF_S2)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S2, + ) + cute.arch.sync_warp() + if lane < TOP_K: # every writer lane arrives (its own release of s_ids / s_w) + prims.mbarrier_arrive(topk_bar.subview(0)) # this token's ids / weights are in smem + if t == cutlass.Int32(0): + prims.mbarrier_arrive( + tok0_bar.subview(0) + ) # token 0's experts: the early group-0 stages + cute.arch.sync_warp() + + if cutlass.const_expr(EARLY0): + if warp == W_WARP: + # slots 0..3 = token 0's four distinct experts; group 0's stage ranges are the same for every G + _wait(tok0_bar.subview(0), 0) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(STAGES): + _issue_fc1( + lane, + n0, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + n0 = n0 + cutlass.Int32(1) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if warp == X_WARP: + # slots: entry i < 4 M is (token i / 4, top-k i % 4); distinct experts in entry order + NCH = (TOP_K * M_MAX + 31) // 32 + _wait(topk_bar.subview(0), 0) + gcnt = cutlass.Int32(0) + rec = [] # per chunk: (first, rank, slot rank of the entry, expert, entry valid) + for ch in cutlass.range_constexpr(NCH): + ei = lane + cutlass.Int32(32 * ch) + valid = ei < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, ei, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + known = cutlass.Boolean(False) + prev = cutlass.Int32(-1) + key = my_e + if cutlass.const_expr(ch > 0): + prev = cutlass.select_( + valid, + s_emap.load(idx=cutlass.select_(valid, my_e, cutlass.Int32(0))), + cutlass.Int32(-1), + ) + known = prev >= cutlass.Int32(0) + key = cutlass.select_(known, cutlass.Int32(-1000) - lane, my_e) + cute.arch.sync_warp() # every lane's s_emap read before any lane's write + same = _match_any(key) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (known == cutlass.Boolean(False)) & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = gcnt + cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + slot_l = cutlass.select_(known, prev, slot_l) + if first: + s_emap.store(rank, idx=my_e) + if valid: + if bx == 0: + route_dbg.store( + my_e, + idx=(ei // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + ei % cutlass.Int32(TOP_K), + ) + rec.append((first, rank, slot_l, my_e, valid)) + cute.arch.sync_warp() + gcnt = gcnt + cute.arch.popc(bal) + # slot of rank r: with a partial last group (tail < 4) and >= 3 groups, that group becomes group 1 (its + # FC2 then is not the last hand-off behind a short FC1 phase); the full groups 1.. shift up by one + ngr = (gcnt + cutlass.Int32(3)) >> cutlass.Int32(2) + tail = gcnt - (ngr - cutlass.Int32(1)) * cutlass.Int32(4) + do_perm = ( + cutlass.Boolean(PERM) + & (tail < cutlass.Int32(4)) + & (ngr >= cutlass.Int32(PERM_MIN_NG)) + ) + for z in cutlass.range_constexpr((G_CAP + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP): + s_el.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + for ch in cutlass.range_constexpr(NCH): + first, rank, slot_r, my_e, valid = rec[ch] + ei = lane + cutlass.Int32(32 * ch) + for which in cutlass.range_constexpr(2): + r = rank if which == 0 else slot_r + sl_m = cutlass.select_( + r < cutlass.Int32(4), + r, + cutlass.select_( + r >= (ngr - cutlass.Int32(1)) * cutlass.Int32(4), + r - (ngr - cutlass.Int32(2)) * cutlass.Int32(4), + r + cutlass.Int32(4), + ), + ) + sl_m = cutlass.select_(do_perm, sl_m, r) + if cutlass.const_expr(which == 0): + if first: + s_el.store(my_e, idx=sl_m) + else: + if valid: + s_kpos.store( + ei % cutlass.Int32(TOP_K), + idx=sl_m * cutlass.Int32(M_MAX) + ei // cutlass.Int32(TOP_K), + ) + # slots per group + if lane < cutlass.Int32(NG_CAP): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + nh = cutlass.select_(do_perm & (lane == cutlass.Int32(1)), tail, nh) + nh = cutlass.select_( + do_perm & (lane == ngr - cutlass.Int32(1)), cutlass.Int32(4), nh + ) + s_gn.store(nh, idx=lane) + cute.arch.sync_warp() + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + prims.mbarrier_arrive( + route_bar.subview(0) + ) # every lane: its slots / kpos / plan-row writes are in smem + cute.arch.sync_warp() + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = s_gn.load(idx=lane) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add(cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad) + else: + _wait(route_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(B1_WARP < 0): + if warp == cutlass.Int32(0): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + if warp < 4: # B1 scales to TMEM (the MMA needs them only with the first ring stage) + _wait(b1_load.subview(0), 0) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + n_iss = cutlass.Int32( + 0 + ) # stages issued (early ones included): the L2-prefetch warp's throttle + w_emp = cutlass.Int64(0) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _dbg_delay(bx) + _issue_fc1( + lane, + stage, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _issue_fc2( + lane, + stage, + g, + blk, + t2, + n_here, + s_el, + sA, + sSF, + ab_full, + tma_a2_desc.get_ptr(), + tma_s2_desc.get_ptr(), + s2w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if lane == 0: + _stamp(tdbg, bx, 3) + _tstore(tdbg, bx, 30, w_emp) + + # ---- (SFA_CP = 0) scale splicer (warp 6): atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if cutlass.const_expr(not SFA_CP): + if warp == S_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt_ = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt_ = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt_ = cutlass.Int32(KT2) + for _j in cutlass.range(cnt_, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32( + 4 + ) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSF.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSF.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSF.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSF.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + # ---- FC2 B loader (warp 5): group g's image once its counter is complete + if warp == X_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for g in range(ng): + buf = g % cutlass.Int32(NB2BUF) + if g >= cutlass.Int32(NB2BUF): + _waitd( + b2_free.subview(buf), + ((g // cutlass.Int32(NB2BUF)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 29) + caddr = cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2_POLL_NS) + cv = _ld_acquire_u32(caddr) + cute.arch.sync_warp() + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 51) + _fence_proxy_async_global() + if g >= cutlass.Int32(1): + _waitd( + sfb2_free.subview(0), (g - cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 107 + ) + if lane == 0: + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(0), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), tdbg, 108 + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 52) + if g >= cutlass.Int32(1): + # the single TMEM scale slot: the previous group's FC2 MMAs are done + gp = g - cutlass.Int32(1) + _waitd( + b2_free.subview(gp % cutlass.Int32(NB2BUF)), + (gp // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), + tdbg, + 109, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 53) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for aa in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFB2_COL + aa * 4), + sd0 + cutlass.Int64(aa * (512 >> 4)), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + prims.tcgen05_commit(sfb2_free.subview(0)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + + # ---- (PFD) L2 prefetch warp (12): the same stage sequence as the producer, PFD stages ahead of it + if cutlass.const_expr(PFD > 0): + if warp == cutlass.Int32(PF_WARP): + pf = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + kx = j * cutlass.Int32(C1 * KT_BYTES) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_prefetch( + tma_a1_desc.get_ptr(), + [ + kx + cutlass.Int32(q * KT_BYTES), + c * cutlass.Int32(R1), + e, + ], + [], + ) + atom0 = ( + e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2)) + ) * cutlass.Int32(SF1_KA) + j * cutlass.Int32(16) + _l2_prefetch( + s1w.subview(atom0 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 8192, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load( + idx=g * cutlass.Int32(4) + cutlass.Int32(jj) + ) + prims.cp_async_bulk_tensor_prefetch( + tma_a2_desc.get_ptr(), + [ + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ], + [], + ) + a2 = ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) * cutlass.Int32(SF2_KA) + cutlass.Int32(4 * t2) + _l2_prefetch( + s2w.subview(a2 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 2048, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + + # ---- (REL_WARP) group releases (warp 8): this CTA's finished units of group p, one release add per phase + if cutlass.const_expr(REL_WARP): + if warp == cutlass.Int32(8): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + grp_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + _dbg_delay(bx * cutlass.Int32(13) + p) + if (lane == 0) & (n_done > cutlass.Int32(0)): + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + _red_release_add( + cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + if (lane == 0) & (p == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 54) + cute.arch.sync_warp() + + # ---- MMA (warp 7) + if warp == MMA_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + m_b2 = cutlass.Int64(0) + m_full = cutlass.Int64(0) + m_acc = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + m_acc = _waitt( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + cute.arch.sync_warp() + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g % cutlass.Int32(NB2BUF) + m_b2_0 = m_b2 + m_b2 = _waitt( + b2_ready.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), m_b2 + ) + if cutlass.const_expr(TIMING): + if (lane == 0) & (g < cutlass.Int32(4)): + tdbg.store( + m_b2 - m_b2_0, + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(44) + g, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + m_acc = _waitt( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB2_COL + cutlass.Int32(t2 * NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + cute.arch.sync_warp() + k2 = k2 + cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit( + b2_free.subview(buf) + ) # this group's B2 tile and TMEM scales are free + cute.arch.sync_warp() + if lane == 0: + _tstore(tdbg, bx, 26, m_b2) + _tstore(tdbg, bx, 27, m_full) + _tstore(tdbg, bx, 28, m_acc) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + img8 = cutlass.Array( + img32.data_ptr(0), shape=(NG_CAP * IMG_BYTES,), dtype=cutlass.Int8, alignment=16 + ) + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if qi < cutlass.Int32(6 * G_CAP): + e_sl = s_el.load(idx=sl_e) + if e_sl >= cutlass.Int32(0): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32(escal.load(idx=fld * cutlass.Int32(E_LOCAL) + e_sl)), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + e_wait = cutlass.Int64(0) + e_busy = cutlass.Int64(0) + e_a = cutlass.Int64(0) + e_pw = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + unit_done = (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)) + e_wait = _waitt( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + tb0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1_LD, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=(t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + tc0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_a = e_a + (tc0 - tb0) + g_u = slot >> cutlass.Int32(2) + jrow = slot & cutlass.Int32(3) + # this warp's tokens t = warp + 4 tr, all in flight together (latency-bound otherwise) + NT = (M_MAX + 3) // 4 + xr = cute.make_rmem_tensor((NT,), cutlass.Float32) + for tr in cutlass.range_constexpr(NT): + tq = cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + xr[tr] = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=(tq * C1 + 0) * R1 + lane), + s_part.load(idx=(tq * C1 + 1) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 2) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 3) * R1 + lane), + ) + if (j0 == cutlass.Int32(0)) & (j1 != cutlass.Int32(S1)): + # head of a unit split with the next CTA: its partial sums + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + _st_relaxed_u32( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint(), + xr[tr].bitcast(cutlass.Int32), + ) + else: + if j0 != cutlass.Int32(0): + # tail: the head's partial sums (head first, as TRTLLM-Gen's K order) + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + pbits = _ld_relaxed_u32(paddr) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xr[tr] = _add_rn(pbits.bitcast(cutlass.Float32), xr[tr]) + cute.arch.sync_warp() + _dbg_delay(bx * cutlass.Int32(7) + warp) + es = slot * cutlass.Int32(8) + ea = s_esc.load(idx=es) + ec = s_esc.load(idx=es + cutlass.Int32(1)) + eal = s_esc.load(idx=es + cutlass.Int32(2)) + ebe = s_esc.load(idx=es + cutlass.Int32(3)) + eli = s_esc.load(idx=es + cutlass.Int32(4)) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + n0 = jrow * cutlass.Int32(M_MAX) + t2 = cblk >> cutlass.Int32(4) + cc = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = g_u * cutlass.Int32(IMG_BYTES) + sfoff = ( + gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2 * cutlass.Int32(NUM_KB) + + ((cblk & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(512) + ) + + (cblk & cutlass.Int32(3)) + ) + for tr in cutlass.range_constexpr(NT): + # (skip) only tokens routed to this slot (warp-uniform); other B2 rows are never used + if s_kpos.load( + idx=slot * cutlass.Int32(M_MAX) + + cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + ) >= cutlass.Int32(0): + xv = xr[tr] + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + qv, sfb = _swiglu_quant( + xv, gate_raw, is_up, ea, ec, eal, ebe, eli + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + t = warp + cutlass.Int32(4 * tr) + if ( + (lane == 0) + & (t < cutlass.Int32(M_MAX)) + & cutlass.Boolean(not DBG_NOSTORE) + ): + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), row n + # at byte 16 (n % 32) + 4 (n / 32), byte c % 4 + n = n0 + t + img32.store( + (lo, hi), + idx=( + gbase + + t2 * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=sfoff + + (n & cutlass.Int32(31)) * cutlass.Int32(16) + + (n >> cutlass.Int32(5)) * cutlass.Int32(4), + ) + tc1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_pw = e_pw + (tc1 - tc0) + if cutlass.const_expr(not REL_WARP): + if unit_done: + if lane == 0: + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + if cutlass.const_expr(EPI_FENCE): + _fence_acq_rel_gpu() + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=EPI_BAR, number_of_threads=EPI_THREADS + ) # s_part reads done + if cutlass.const_expr(not REL_WARP): + if unit_done & (tidx == 0): + _red_release_add( + cnt.subview(g_u * cutlass.Int32(8)).data_ptr().toint(), + cutlass.Int32(1), + ) + if cutlass.const_expr(TIMING): + e_busy = e_busy + (_gtime() - tb0) + k1 = k1 + cutlass.Int32(1) + if cutlass.const_expr(REL_WARP): + prims.mbarrier_arrive( + grp_done.subview(p & cutlass.Int32(3)) + ) # this CTA's group-p units are out + if tidx == 0: + _stamp(tdbg, bx, 10) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + e_wait = _waitt( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + if cutlass.const_expr(TIMING): + if ( + (tidx == 0) + & (g == ng - cutlass.Int32(1)) + & (bb == cutlass.Int32(0)) + ): + _stamp(tdbg, bx, 8) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = [] + ya = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), + cutlass.Float32, + ), + num=ACC2_STRIDE, + ) + for t in cutlass.range_constexpr(M_MAX): + v = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + v = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(ya[w * M_MAX + t]), v + ) + yv.append(v) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if warp < s_gn.load(idx=g): + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + cutlass.BFloat16( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)) + ), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + if tidx == 0: + _stamp(tdbg, bx, 6) + _tstore(tdbg, bx, 31, e_wait) + _tstore(tdbg, bx, 34, e_busy) + _tstore(tdbg, bx, 36, e_a) + _tstore(tdbg, bx, 38, e_pw) + prims.mbarrier_arrive(ys_done.subview(0)) # this thread's ys values are in smem + + # ---- finalize on all 16 warps: (block bb, token t) pairs, the fma chain over t's top-4 in order, + y_sh + NPAIR = NB2_CAP * M_MAX + NPW = (NPAIR + NWARPS - 1) // NWARPS + wyl = [] + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + ok = (pair < cutlass.Int32(NPAIR)) & (bb < nb) + ch = cutlass.select_( + ok, + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wyl.append( + _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + if cutlass.const_expr(SHARED) + else cutlass.Int32(0) + ) + _wait(ys_done.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 19) + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if (pair < cutlass.Int32(NPAIR)) & (bb < nb): + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ), + acc_o, + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SHARED): + wy = wyl[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + bias: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + img32: cute.Tensor, + cnt: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + w1b: cute.Tensor, + w2b: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + bias, + escal, + plan, + partials, + img32, + cnt, + s1w, + s2w, + w1b, + w2b, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..81122a06c18c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3914 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP4 rank (routed and shared intermediate 768): +``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 # capacity: fixed: the checkpoint's hidden size +# routed intermediate per rank: 1536 at TP2 (this target: FC2 K = 6 full 256-K tiles), 768 at TP4, 384 at TP8 (FC2 K +# = 1.5 256-K tiles: the last tile's A box is half out of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 +# MMAs / scale atoms are issued; no padding, no extra bytes) +I_TP = _cfg("i_tp", 768) +TWO_I = 2 * I_TP +assert I_TP % 128 == 0 and 128 <= I_TP <= 1536, I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 # capacity: fixed: the CTA size +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is the low-M MXFP8 GEMM's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# TAIL_TRIGGER (with LATE_TRIGGER, no all-reduce push): the CTAs that own more FC2 blocks than the rest (the first +# H / 32 mod NCTA, which end last) release the dependents right after the grid dependency, so the dependents launch +# when the other CTAs tear down, onto the SMs they leave, instead of after the last of the heavy CTAs. +TAIL_TRIGGER = bool(_cfg("tail_trigger", 1)) +NHEAVY = (H // 32) % NCTA if NCTA < H // 32 else 0 +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +# the shared intermediate per rank: 1536 (TP2), 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate +# positions: its gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP2: 1 x (64 + +# 64), full K in one accumulator, the order of the stock tcgen05 MXFP8 GEMM; TP4: 2 x (32 + 32), the order of +# m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. +SH_I = _cfg("sh_i", 768) +SH_ROWS = SH_I // 24 +# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) +SH_SPLIT = 128 // (2 * SH_ROWS) +SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N +SH_KB = H // 32 # capacity: fixed: a kernel tile constant (gate_up K blocks) (gate_up K blocks) +SH_DKB = SH_I // 32 # down K blocks +SH_DKT = SH_I // 128 # down k-tiles +SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the +# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction +# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) +) # (in-chain: a win at M <= 2, not 4..8) +if SHARED: + assert SH_I in (384, 768, 1536), SH_I + assert SH_PRETILE or SH_SPLIT == 2, ( + "the un-pretiled shared stream is the 32-row (TP4) layout only" + ) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # 6 stages per unit # capacity: fixed: pipeline stages +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) +# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) +KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 + + +def _nkb2(t2: int) -> int: + """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" + return NUM_KB if t2 < KT2 - 1 else KB2_LAST + + +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy +# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few +# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. +HB_REPL = _cfg("hb_repl", 1) +assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) +HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +# Hang guard: every cross-CTA spin (the kernel needs all its CTAs resident at once) traps after this long, so a lost +# co-residency or a missing producer surfaces as a device exception instead of a silent hang. The timer is read only +# inside a spin iteration; the arithmetic is unchanged. +HANG_NS = _cfg("hang_ms", 2000) * 1000000 +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # capacity: fixed: a warp role index +S_WARP = 6 # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +# E0_M4: at M = 4 the MMA warp routes +# token 3 (NRW0 = 4), so the early group-0 stages go to FC2 B builder warp 8 instead (idle until the routing is done; +# RW_EXTRA = 0 at M = 4). Issue order only: exact. +E0_M4 = ( + bool(_cfg("e0_m4", 1)) and NRW0 == 4 and RW_EXTRA == 0 and THREADS // 32 > 8 +) # M == 4 only: at M > 4 warp 8 routes (RW_EXTRA) +E0_WARP = ( + 8 if E0_M4 else MMA_WARP +) # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and (NRW0 <= 3 or E0_M4) and bool(_cfg("e0_off", 1)) +# B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing +# warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the +# two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. +if B1_AFTER_TOPK and XSTAGE and EARLY0 and not E0_OFFLOAD: + raise ValueError( + "m3_moe: b1_after_topk=1 with xstage, early0 and no early-stage offload deadlocks (warp 4 waits on x_bar[1] " + "before TOPK2_BAR, the B1 build waits on TOPK2_BAR before x_bar[1]); set b1_after_topk=0 or early0=0" + ) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# TP2 at 8 tokens: the 98 KB of hidden states fit no 5-stage ring; the router GEMV then reads them from global (the +# kernel's non-XSTAGE path, same math) +if XSTAGE and M_MAX * H * 2 > STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA): + XSTAGE = False +assert not XSTAGE or M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 +_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 +_SH_FIX = ( + SH_GT * SH_NB * 128 + + SH_DKT * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + _SH_XCH_BYTES + + 2 * SH_ROWS * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * SH_NB * 128 +_SB_XC = _SB_BH + SH_DKT * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + _SH_XCH_BYTES +_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = SH_NB +SH_SFA = SH_ACC_DN + 16 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +# pre-tiled stages beyond the ring that warp 0 prefetches into L2, per lane (TP2: 72 stages) +SH_PF = max(2, (SH_STAGES - SH_NST + 31) // 32) if SHARED else 2 +# TP2 (SH_GT 48 gate_up k-tiles, SH_DKT 12) needs more than 256 columns: allocate the next power of two +SH_TMEM = 256 if SH_SFB_DN + 4 * SH_DKT <= 256 else 512 +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 32 * SH_PF, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 4 * SH_DKT <= SH_TMEM, ( + SH_NST, + _SB_END, + BIG, + ) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _trap_after(elapsed, *, loc=None, ip=None): + """trap when elapsed (ns, int64) exceeds HANG_NS.""" + _llvm.inline_asm( + None, + [cutlass.Int64(elapsed).ir_value(loc=loc, ip=ip)], + f"{{ .reg .pred p; setp.gt.s64 p, $0, {HANG_NS}; @p trap; }}", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _hang_tick(t0): + """One spin iteration of a cross-CTA wait: the wait's start (latched on the first call; t0 = 0 before it), after + trapping once the wait has lasted longer than HANG_NS.""" + now = _gtime() + t1 = cutlass.select_(t0 == cutlass.Int64(0), now, t0) + _trap_after(now - t1) + return t1 + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + # (a partial last k-tile has no atoms past I_TP) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(SH_PF): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(SH_PF) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SH_NB, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if cutlass.const_expr(SH_SPLIT == 1): + # TP2: one K slice; quadrant qd = rows 32 qd .. (gate positions 0..63, then up 0..63), no partials: + # the owner sum is (0 + acc), then bf16, at _SB_GU word (u SH_ROWS + p) 8 + t = (32 qd + l) 8 + t + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + elif cutlass.const_expr(SH_SPLIT == 2): + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + else: + # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are + # accumulator columns 8 qd + t + av = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_GU + qd * cutlass.Int32(8), + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(1): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd == cutlass.Int32(0): + # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + for s_ in cutlass.range_constexpr(1, SH_SPLIT): + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + lane * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(SH_ROWS * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + t0h = cutlass.Int64(0) + while pend != cutlass.Int32(0): + t0h = _hang_tick(t0h) + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(SH_DKT): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(LATE_TRIGGER and TAIL_TRIGGER and NHEAVY > 0): + if ar_off & (bx < cutlass.Int32(NHEAVY)): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + t0h = cutlass.Int64(0) + while done == cutlass.Int32(0): + t0h = _hang_tick(t0h) + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + t0h = cutlass.Int64(0) + while (cv - cnt_target) < cutlass.Int32(0): + t0h = _hang_tick(t0h) + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + for rr in cutlass.range_constexpr(HB_REPL): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4) + + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + hb_off = cutlass.Int32(0) + if cutlass.const_expr(HB_REPL > 1): + # this CTA's copy of the records + hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + t0h = cutlass.Int64(0) + while done == cutlass.Int32(0): + t0h = _hang_tick(t0h) + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview( + (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) + ), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + t0h = cutlass.Int64(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + t0h = _hang_tick(t0h) + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if cutlass.const_expr(HB_REPL == 1): + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + else: + if lane < HB_REPL: # copy `lane` + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4) + + lane * cutlass.Int32(HB_STRIDE), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + t0h = cutlass.Int64(0) + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + t0h = _hang_tick(t0h) + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op.py new file mode 100644 index 000000000000..ab60e7d96370 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op.py @@ -0,0 +1,779 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP4 rank (routed and shared intermediate 768 per rank), +over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm.logger import logger + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +INTER = 768 # capacity: fixed: the routed experts' width per tp4 rank +NUM_EXPERTS = 128 +SHARED_INTER = 768 # capacity: fixed: the shared expert's width per tp4 rank +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size (m3_moe's shared-expert CTAs (beside num_ctas routed CTAs)) +TOP_K = 4 +U1 = 2 * INTER // 32 # FC1 units per expert +# The router GEMV takes at most 2 experts per CTA (router weights held in registers). +MIN_CTAS = 64 # capacity: fixed: an SM split, not a size + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # FC1 stages (items) per unit # capacity: fixed: pipeline stages +_ITEM_BYTES = _UNIT_BYTES // _S1 # capacity: fixed: a ring item's bytes +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters. Off at TP4 (at 4 / 8 tokens +# 0.5 / 0.8 us/call slower than the tagged records in the TP4 sweep). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 0) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit). 5 at 8 tokens (with a 50 ns FC2 record poll below) +# measured 0.8 us/call faster than 6 in the TP4 sweep (same MMAs, bit-identical outputs); at 4 tokens the default is +# fastest with heavy_fc2_only. +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 5) +# The FC2 B builders' record poll back-off in ns (0 = the kernel's default, 200); 50 at 8 tokens. +B2_POLL_NS_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 50) +# ``early0``: warp 4 streams group 0 right after its own top-k. Off at 8 tokens (0.6 us/layer faster in the chain: +# there it is on every routing warp's critical path and the ring holds only 5 stages). +EARLY0_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 0) +# The plan's ``heavy_fc2_only``: 192 FC2 output blocks over 188 routed CTAs leave 4 CTAs with two blocks, whose +# trailing FC2 after the last hand-off otherwise ends the kernel 3-7 us after the rest; they stream no FC1 instead +# (1.7 / 0.9 us/call faster at 4 / 8 tokens in the TP4 sweep; bit-identical outputs). +HEAVY_FC2_ONLY_BY_M: Tuple[int, ...] = (0, 0, 0, 1, 0, 0, 0, 1) +# ``hb_repl``: copies of the tagged FC1 -> FC2 records (applied when > 1); CTA bx polls copy bx % hb_repl, which +# spreads the polls of the few hot record lines over L2. 7 and 8 tokens hand off through the B images instead. +HB_REPL_BY_M: Tuple[int, ...] = ( + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, +) # off at TP4 (0.2 us/call slower at 4 tokens) +# ``lg_repl``: replicas of the router logit exchange at up to LG_REPL_MAX_TOKENS tokens (fewer stores before the +# exchange; bit-identical outputs). Above that the kernel's default applies. +LG_REPL = 32 +LG_REPL_MAX_TOKENS = 0 # capacity: fixed: measured crossover (off at TP4: flat in the sweep) +# The TP4 FLAT stack: the TP2 FLAT values (no B2 images, 16 router-logit replicas, AR push parities loaded in the +# finalize, PDL release right after the grid wait) over the per-m tables above. Bitwise-identical output (m 1..8 x 4 +# routings, eager + graph); standalone m=4 -0.49 us/call (late_trigger). The values are part of the engine's module +# key. +FLAT_STACK = {"b2img": 0, "lg_repl": 16, "ar_pre": 0, "late_trigger": 0} +# FLAT only up to 4 tokens: at 8 tokens the per-m tables are faster (TP2: 2.96 us/call; TP4 in-model at 2 requests +# +175 us/step with FLAT). Bit-identical either way. +FLAT_MAX_TOKENS = 4 # capacity: fixed: measured crossover + + +def flat_applies(num_tokens: int) -> bool: + """Whether an engine state for ``num_tokens`` tokens runs the FLAT stack.""" + return num_tokens <= FLAT_MAX_TOKENS + + +# Byte-balanced FC1 ranges below one unit per CTA (:func:`_fine_ranges`) for groups too small to give every CTA a +# whole unit; False gives such groups whole units on the CTAs with the least FC2 work. +FINE_PLAN = True +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 + = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row + R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then + down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per + layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + P = I_ // n_cta + split = 128 // (2 * P) + gt = H_ // 128 // split + dkt = I_ // 128 + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + octant = r // P + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(2 * dkt): + jt, kt = gd // dkt, gd % dkt + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, gt + gd, :16384] = img.reshape(-1) + out[s, gt + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def _fine_ranges(weights, total: int): + """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item + counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary + strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit + runs at least to that unit's end. Returns the item counts.""" + n = len(weights) + items = [0] * n + s = 0 + acc = 0.0 + for c in range(n): + acc += weights[c] + e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) + if s % _S1 and e < total: # started inside a unit: run to (at least) its end + e = max(e, min(total, (s // _S1 + 1) * _S1)) + items[c] = e - s + s = e + assert s == total + return items + + +def make_phase_plan( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = FINE_PLAN, + heavy_fc2_only: bool = False, +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + # (heavy_fc2_only) CTAs with more FC2 blocks than the minimum stream no FC1: they run each group's FC2 as soon as + # its hand-off lands, so their trailing FC2 after the last hand-off matches the others' two lagged groups + min_nb = min(nbs) + no_fc1 = [heavy_fc2_only and nbs[c] > min_nb for c in range(num_ctas)] + BIG_W = 1 << 40 + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + for c in range(num_ctas): + if no_fc1[c]: + fc2[c] += BIG_W + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [ + 0 if no_fc1[c] else max(_S1, (limit - f) // _ITEM_BYTES) + for c, f in enumerate(fc2) + ] + + lo, hi = 0, max(f for c, f in enumerate(fc2) if not no_fc1[c]) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 + lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] + sw = sum(wts) + items = _fine_ranges([w * total / sw for w in wts], total) + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + inter: int = INTER, + ): + """``inter``: the routed intermediate per rank; this engine is built for INTER (768) only.""" + if inter != INTER: + raise ValueError( + f"m3_moe (TP4): routed intermediate {inter} per rank is not supported (only {INTER})" + ) + self.inter = inter + logger.info(f"[M3 MoE] Rubin m3_moe engine (tp4 tables): state for {num_tokens} tokens") + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + if self.num_ctas > HIDDEN // 32: + # a routed CTA that owns no 32-row FC2 output block deadlocks the FC2 hand-off on its second call + raise ValueError( + f"m3_moe runs at most {HIDDEN // 32} routed CTAs (one FC2 output block each), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, + num_tokens, + M3_MOE_LAG, + 0.0 if num_tokens == 1 else 1.0, + True, + heavy_fc2_only=bool(HEAVY_FC2_ONLY_BY_M[num_tokens - 1]), + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "i_tp": INTER, + } + if HB_REPL_BY_M[num_tokens - 1] > 1: + self.config["hb_repl"] = HB_REPL_BY_M[num_tokens - 1] + if num_tokens <= LG_REPL_MAX_TOKENS: + self.config["lg_repl"] = LG_REPL + if shared: + self.config["sh_i"] = SHARED_INTER + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + if B2_POLL_NS_BY_M[num_tokens - 1] > 0: + self.config["b2_poll_ns"] = B2_POLL_NS_BY_M[num_tokens - 1] + if not EARLY0_BY_M[num_tokens - 1]: + self.config["early0"] = 0 + if flat_applies(num_tokens): + self.config.update(FLAT_STACK) + # engagement line (every rank; once per engine state = once per token count) + from tensorrt_llm.logger import logger as _flog + + _flog.info( + "MiniMax-M3 MoE engine FLAT (TP4): " + + " ".join(f"{k}={self.config[k]}" for k in FLAT_STACK) + + f" tokens={num_tokens}" + ) + else: + from tensorrt_llm.logger import logger as _flog + + _flog.info( + f"MiniMax-M3 MoE engine FLAT (TP4) off above {FLAT_MAX_TOKENS} tokens (per-m tables): " + f"tokens={num_tokens}" + ) + # E0@m4: at 4 tokens the early group-0 FC1 stages are issued by FC2 B builder warp 8 instead of the routing + # warp 4 (issue order only, bitwise; TP4 -0.19 / TP8 -0.21 us/call). + if num_tokens == 4: + from tensorrt_llm.logger import logger as _e0log + + _e0log.info( + "MiniMax-M3 MoE engine E0M4 (%s): early group-0 stages on warp 8, tokens=4" % "TP4" + ) + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + # Every CTA waits on others (cross-CTA spins on partials, hand-off records, the shared expert's h and y): the + # grid must be resident at once, one CTA per SM. The module's shared-memory / TMEM budget assertion makes one + # CTA fit an SM; the grid must not exceed the SMs. (The kernel also traps after HANG_NS in any such spin.) + if self.num_ctas + mod.NSH > sms: + raise ValueError( + f"m3_moe: a grid of {self.num_ctas} + {mod.NSH} CTAs exceeds the {sms} SMs; its CTAs wait on each " + "other and must all be resident at once" + ) + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full( + (mod.G_CAP * mod.U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw + ) + # the FC1 -> FC2 records; with HB_REPL > 1 copies, HB_STRIDE int32 words apart + hb_words = mod.G_CAP * num_tokens * mod.U1 * 4 + if mod.HB_REPL > 1: + hb_words = mod.HB_REPL * mod.HB_STRIDE + self.hbuf32 = torch.zeros(hb_words, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [3072, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 768], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + # keyed on the config the kernel module was built with (not self.config), so a module built with any other + # trace-time config can never reuse another build's compiled kernel + key = (_ENGINE_PATH,) + tuple(sorted(self.mod.M3_CONFIG.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op16.py new file mode 100644 index 000000000000..446776419315 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op16.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host side of ``m3_moe16`` (m3_moe16_kernel.py): the MiniMax-M3 MoE layer (fused router + routed experts + shared +expert) for 9..16 decode tokens on one TP4 rank. The 1..8-token engine (op.M3MoeState) is untouched; select this +state only for num_tokens > 8. + +Output: bf16 [M, 6144] = bf16(shared + routed), the routed part bit for bit TRTLLM-Gen's fused MoE on the fp32 router +GEMV logits (m3_router_gemv's math), the shared part as the in-model m3_mx_gemv shared expert. + + state = M3Moe16State(device, num_tokens=12) # before CUDA-graph capture; one per token count + layer = state.layer(w3_w1_weight, w3_w1_weight_scale, w2_weight, w2_weight_scale, fc31_alpha, fc31_scale_c, + act_alpha, act_beta, act_clamp, fc2_alpha, e_score_correction_bias, + router_weight=router_weight, fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf)) + out = layer(hidden) # first call compiles: call once before capture + +Protocol state (per state object, shared by the layers that run on it in stream order): per-group FC2-B images and +unit counters, the B1 image + counter, per-CTA epochs. Every call adds exactly 192 to each group counter and NCTA to +the B1 counter, so nothing is reset between calls (graph replay safe). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch + +from .op import ( + HIDDEN, + SHARED_CTAS, + SHARED_INTER, + U1, + M3MoeLayer, + _compiled, + _kernel_module, + _shared_views, + _view, + is_supported, +) + + +def _group_slots(g: int, p: int, perm: bool, min_ng: int = 5) -> int: + """Slots of group p when the call has g distinct experts (the kernel's slot order).""" + ng = (g + 3) // 4 + tail = g - 4 * (ng - 1) + if p >= ng: + return 0 + if perm and tail < 4 and ng >= max(3, min_ng): + return tail if p == 1 else 4 + return min(4, g - 4 * p) + + +def make_phase_plan16(num_ctas: int, m_max: int, lag: int = 2, perm: bool = True, min_ng: int = 5): + """op.make_phase_plan (early0, tail_w 1) with the kernel's group sizes (PERM: a partial last group is group 1).""" + S1 = 6 + UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) + ITEM = UNIT_BYTES // S1 + blk_bytes = 32 * (768 // 2 + 768 // 16) + gcap = 4 * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + for g in range(4, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = _group_slots(g, p, perm, min_ng) + total = n_p * U1 * S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * _group_slots(g, p - lag, perm, min_ng) * blk_bytes + if p == ng - 1 and not (p == 0 and m_max >= 2): + for q in range(max(0, ng - lag), ng): + fc2[c] += nbs[c] * _group_slots(g, q, perm, min_ng) * blk_bytes + items = [0] * num_ctas + if total >= S1 * num_ctas: + + def counts_for(limit): + return [max(S1, (limit - f) // ITEM) for f in fc2] + + lo, hi = 0, max(fc2) + total * ITEM + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > S1), + key=lambda c: (-(fc2[c] + items[c] * ITEM), c), + ) + items[c] -= 1 + extra -= 1 + else: + for c in sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // S1]): + items[c] = S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +_ENGINE16_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe16_kernel.py") +M16_MIN_TOKENS, M16_MAX_TOKENS = 9, 16 +# Default routed CTAs (beside the SHARED_CTAS shared-expert CTAs). +M16_NUM_CTAS = 128 +# Phase lag: phase p streams FC1 of group p and FC2 of group p - M16_LAG (the kernel's default lag). +M16_LAG = 3 + + +class M3Moe16State: + """Workspace + compiled kernel of ``m3_moe16`` for ``num_tokens`` (9..16) tokens per call.""" + + fused_router = True + shared = True + + def __init__(self, device: torch.device, num_tokens: int, num_ctas: Optional[int] = None): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3Moe16State allocates and compiles: build it before CUDA-graph capture" + ) + if not M16_MIN_TOKENS <= num_tokens <= M16_MAX_TOKENS: + raise ValueError( + f"m3_moe16 takes {M16_MIN_TOKENS}..{M16_MAX_TOKENS} tokens per call, not {num_tokens}" + ) + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe16: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or M16_NUM_CTAS + if not 128 <= self.num_ctas <= sms - SHARED_CTAS: + raise ValueError( + f"m3_moe16 runs 128 .. {sms - SHARED_CTAS} routed CTAs + {SHARED_CTAS} shared (all resident)" + ) + self.m = num_tokens + plan, nb_cap = make_phase_plan16(self.num_ctas, num_tokens, M16_LAG) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "shared": 1, + } + try: + self.mod = mod = _kernel_module(self.config, _ENGINE16_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe16: num_tokens={num_tokens} does not fit the smem / TMEM budget: {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + # split units' head sums: every slot empty (0xFFFFFFFF) between calls + self.partials = torch.full((mod.G_CAP * U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw) + # FC2-B images per group + the B1 image; counters: group g at 8 g, B1 at 8 NG_CAP, CTA tickets at 8 NG_CAP + 8 + self.img = torch.zeros(mod.IMG_TOTAL // 4, dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 16, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros( + mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw + ) # logits: tag << 32 | f32 + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> M3MoeLayer: + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + layer.shared_weights = _shared_views(shared_weights, "m3_moe16") + layer.shared_mxfp8 = tuple(shared_weights) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError(f"m3_moe16: expects contiguous bf16 hidden states [{m}, {HIDDEN}]") + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + layer.bias, + layer.escal, + self.plan.view(-1), + self.partials, + self.img, + self.cnt, + layer.w3_w1_weight_scale.reshape(-1).view(torch.int32), + layer.w2_weight_scale.reshape(-1).view(torch.int32), + layer.w3_w1_weight.reshape(-1).view(torch.int8), + layer.w2_weight.reshape(-1).view(torch.int8), + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + *layer.shared_weights, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_moe16 compiles on its first call: call it once before capture" + ) + import cutlass.cute as cute + + sig = [_view(t, 16 if t.element_size() * t.numel() % 16 == 0 else 4) for t in args] + key = (_ENGINE16_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp4/kernels/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_report.py new file mode 100644 index 000000000000..7f864bd240c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_report.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast-path reports a kernel module makes from inside a call: ``taken`` (an engagement line), ``missed`` (a covered +shape that cannot run: raises under ``modeling_v2='require'``), ``uncovered`` and ``capture_missed``. + +The policy (logging once, raising under require) is the target's: ``modeling.py`` installs its functions here at import, +so a kernel reports without importing the target module. Until then every report is a no-op. +""" + +from __future__ import annotations + +from typing import Callable, Optional + + +def _noop(*args, **kwargs) -> None: + return None + + +taken: Callable[..., None] = _noop +missed: Callable[..., None] = _noop +uncovered: Callable[..., None] = _noop +capture_missed: Callable[..., None] = _noop + + +def install( + *, + on_taken: Callable[[str, Optional[int], str], None], + on_missed: Callable[[str, Optional[int], str], None], + on_uncovered: Callable[[str, str], None], + on_capture_missed: Callable[[str, Optional[int], str], None], +) -> None: + """Route this module's reports to the target's policy.""" + global taken, missed, uncovered, capture_missed + taken, missed, uncovered, capture_missed = on_taken, on_missed, on_uncovered, on_capture_missed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_sampled_accept.py new file mode 100644 index 000000000000..dbeb5c8850a6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_sampled_accept.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). + +Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +active must stay on the generic path; the caller decides. + +Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +(the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens heavier +than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the nucleus +with probability >= top_p, so most rows decide on their second exchange. + +Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based Philox +keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one 4-word +record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and reads +the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the same +order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing and +two call halves keep a rank one call ahead from overwriting a slot still being read. +""" + +from typing import Optional + +import torch + +from ..ops.spec.m3_sampled_accept import m3_sampled_accept as _sampled_accept_op +from ..ops.spec.m3_sampled_accept_ws_words import m3_sampled_accept_ws_words as _ws_words_op + +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu: a row's CTAs +# launch as one thread-block cluster, bounded record polls) and reached through the ops/spec entry m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) + + +class SampledAcceptWorkspace: + """Exchange slots behind one multicast mapping (two call halves), the (finished rows, call) counters, the per-row + token scratch, an extra-attempt counter and the SA-1R round-0 decided counter. Created collectively, outside + CUDA-graph capture.""" + + def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): + self.uc, self.mc_ptr, self.rank, self.world = uc, mc_ptr, rank, world + self.handle, self.comm = handle, comm + dev = uc.device + self.counters = torch.zeros(2, dtype=torch.int32, device=dev) + self.tokens = torch.zeros(MAX_ROWS, dtype=torch.int32, device=dev) + self.extra_attempts = torch.zeros(1, dtype=torch.int32, device=dev) + self.decided = torch.zeros(1, dtype=torch.int32, device=dev) + + @classmethod + def create(cls, mapping, device, one_round_k: int) -> "SampledAcceptWorkspace": + # SA-1R's records carry the round-0 list words: the size depends on the variant (``one_round_k``). + words = int(_ws_words_op(max(mapping.tp_size, 1), one_round_k)) + if mapping.tp_size == 1: + uc = torch.zeros(words, dtype=torch.int64, device=device) + return cls(uc, uc.data_ptr(), 0, 1) + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 8, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("SampledAcceptWorkspace.create failed on another rank") + return cls(uc, mc.data_ptr(), mapping.tp_rank, mapping.tp_size, handle, comm) + + +_WS: Optional[SampledAcceptWorkspace] = None + + +def workspace(mapping, device, one_round_k: int) -> Optional[SampledAcceptWorkspace]: + """The process-wide workspace for the SA-1R variant ``one_round_k`` (the same on every call of a process).""" + global _WS + if _WS is None: + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = SampledAcceptWorkspace.create(mapping, device, one_round_k) + return _WS + + +def sampled_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: Optional[torch.Tensor], + seeds: torch.Tensor, + offsets: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, + ws: Optional[SampledAcceptWorkspace] = None, + use_pdl: bool = True, + *, + one_round_k: int, + prewait_params: bool, + early_trigger: bool, + ilp_walk: bool, +): + """Sampled verification (temperature + top-p) from this rank's bf16 logits shard [nreq * (draft_len + 1), V / tp]: + outputs as shard_accept (accepted [nreq, D+1], num_accepted [nreq], gather_ids [nreq]); None + when unavailable. Per-row sampling parameters and Philox (seed, offset) are the spec metadata's, one per row. + top_ks / min_ps are not applied: rows with top-k or min-p active must be routed to the generic path by the + caller. ``one_round_k`` (SA-1R: 0, 9 or 82), ``prewait_params``, ``early_trigger`` and ``ilp_walk`` select the + kernel variant (the target's constants).""" + del top_ks, min_ps + nreq, draft_len = draft_tokens.shape + rows = nreq * (draft_len + 1) + if logits.dim() != 2 or logits.shape[0] != rows or rows > MAX_ROWS or logits.shape[1] > 51200: + return None + if logits.dtype != torch.bfloat16 or logits.stride(1) != 1: + return None + ws = ws if ws is not None else workspace(mapping, logits.device, one_round_k) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + _sampled_accept_op( + logits, + int(vocab_offset), + temperatures[:rows].float().contiguous(), + top_ps[:rows].float().contiguous(), + seeds[:rows].long().contiguous(), + offsets[:rows].long().contiguous(), + ws.uc, + ws.mc_ptr, + ws.counters, + ws.tokens, + ws.rank, + ws.world, + draft_tokens.int().contiguous(), + batch_idx.int().contiguous(), + accepted, + num_accepted, + gather_ids, + int(force_total), + next_new, + force_vec, + ws.extra_attempts, + bool(use_pdl), + prewait_params=bool(prewait_params), + early_trigger=bool(early_trigger), + one_round_k=int(one_round_k), + decided=ws.decided, + ilp_walk=bool(ilp_walk), + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_select_long_dsl.py new file mode 100644 index 000000000000..49bcef462239 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_select_long_dsl.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). + +Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp4_select_long``, +one 1024-thread CTA per row), from a CTA of 384 threads: + +* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and + mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's + valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned + key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped + block and masks the result, so a thread's loads issue back to back. +* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of + the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. +* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by + that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate + against the whole row instead. +* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. + +The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait +(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are +compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and +loads them before the executor's warmup, so a decode step never compiles one. +""" + +from typing import Dict, Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.utils import ( + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +SELECT_THREADS = 384 +SELECT_WARPS = SELECT_THREADS // 32 +# 1M tokens + draft tokens at 128 tokens per block. +MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row +MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging +TOPK = 16 # capacity: fixed: MSA's top-k +NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) +# The key of -inf (the reference's 0xff800000): never emitted. +NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 +# (keys per thread, rows up to this many valid blocks) +ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) + +# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare +# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group +# minima [16], per-warp candidate counts [SELECT_WARPS]. +_CAND = 0 +_PICK_KEY = _CAND + 2 * MAX_CAND +_PICK_BLOCK = _PICK_KEY + TOPK +_WARP_MIN = _PICK_BLOCK + TOPK +_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 +_WARP_CNT = _GROUP_MIN + 16 +SMEM_WORDS = _WARP_CNT + SELECT_WARPS +SMEM_BYTES = 4 * SMEM_WORDS + + +def _key(score): + """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" + bits = cutlass.Float32(score).bitcast(cutlass.Int32) + return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) + + +def _imin(a, b): + return cutlass.select_(b < a, b, a) + + +def _imax(a, b): + return cutlass.select_(b > a, b, a) + + +def _forced_key(scores, head, block, query, init_blocks, local_start): + raw = scores[head, block, query] + s = cutlass.select_( + block >= local_start, + cutlass.Float32(LOCAL_SCORE), + cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), + ) + return _key(s) + + +@cute.jit +def _select_row( + row, + scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides + n_valid: cute.Tensor, # int32 [>= total_q] + out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] + init_blocks, + local_blocks, + num_kv_heads, + smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory + bar_id: cutlass.Constexpr, + loaded_trigger: cutlass.Constexpr, + raw_prewait, + n_valid_prewait: cutlass.Constexpr, +): + """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier + ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the + row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by + the caller before its grid wait; otherwise it is read here.""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % cutlass.Int32(32) + warp = tid // cutlass.Int32(32) + query = row // num_kv_heads + head = row % num_kv_heads + if cutlass.const_expr(n_valid_prewait): + raw = raw_prewait + else: + raw = n_valid[query] + valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) + local_start = _imax(raw - local_blocks, cutlass.Int32(0)) + cand = cute.make_tensor( + cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) + ) + + if tid < TOPK: + smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) + smem[_PICK_BLOCK + tid] = cutlass.Int32(0) + + if valid <= TOPK: + # Every valid block is picked (thread t holds block t). + if tid < valid: + smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) + smem[_PICK_BLOCK + tid] = tid + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): + items = ITEM_CLASSES[ci][0] + lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] + hi = ITEM_CLASSES[ci][1] + if (valid > lo) & (valid <= hi): + key = cute.make_rmem_tensor(items, cutlass.Int32) + local_min = cutlass.Int32(NO_KEY) + for s in cutlass.range_constexpr(items): + # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. + block = tid + cutlass.Int32(s * SELECT_THREADS) + k = _forced_key( + scores, head, _imin(block, valid - 1), query, init_blocks, local_start + ) + key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) + local_min = _imin(local_min, key[s]) + local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) + if lane < 16: + smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + if tid < 16: + g = cutlass.Int32(NO_KEY) + for w in cutlass.range_constexpr(SELECT_WARPS): + g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) + smem[_GROUP_MIN + tid] = g + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + bound = cutlass.Int32( + -2147483648 + ) # the smallest key: signed keys, unlike the reference's unsigned ones + for g in cutlass.range_constexpr(16): + bound = _imax(bound, smem[_GROUP_MIN + g]) + + # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(items): + cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) + incl = cnt + for d in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync( + incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) + ) + incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) + if lane == 31: + smem[_WARP_CNT + warp] = incl + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + before = cutlass.Int32(0) + total = cutlass.Int32(0) + for w in cutlass.range_constexpr(SELECT_WARPS): + c = smem[_WARP_CNT + w] + before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) + total += c + slot = before + incl - cnt + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + if slot < MAX_CAND: + cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( + tid + cutlass.Int32(s * SELECT_THREADS) + ) + slot += 1 + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if total <= MAX_CAND: + for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): + mine = cand[i] + rank = cutlass.Int32(0) + for j in cutlass.range(total, unroll=8): + rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) + smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) + else: + # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + block = tid + cutlass.Int32(s * SELECT_THREADS) + rank = cutlass.Int32(0) + for j in cutlass.range(valid, unroll=1): + kj = _forced_key(scores, head, j, query, init_blocks, local_start) + less = (kj < key[s]) | ((kj == key[s]) & (j < block)) + rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = key[s] + smem[_PICK_BLOCK + rank] = block + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if warp == 0: + # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. + v = cutlass.Int32(NO_KEY) + if lane < TOPK: + pk = smem[_PICK_KEY + lane] + if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): + v = smem[_PICK_BLOCK + lane] + for size_log in cutlass.range_constexpr(1, 6): + for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): + other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) + ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) + lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) + v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) + if lane < TOPK: + out[query, head, lane] = cutlass.select_( + v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v + ) + + +class _SelectLongDslKernel: + """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count + before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because + ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's + first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it + (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" + + def __init__( + self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool + ): + self.use_pdl = use_pdl + self.early_trigger = early_trigger + self.loaded_trigger = (not early_trigger) and trigger_phase == 1 + self.n_valid_prewait = n_valid_prewait + + @cute.kernel + def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): + row, _, _ = cute.arch.block_idx() + smem = SmemAllocator() + words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) + raw_prewait = cutlass.Int32(0) + if cutlass.const_expr(self.n_valid_prewait): + raw_prewait = n_valid[row // num_kv_heads] + griddepcontrol_wait() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + _select_row( + row, + scores, + n_valid, + out, + init_blocks, + local_blocks, + num_kv_heads, + words, + bar_id=1, + loaded_trigger=self.loaded_trigger, + raw_prewait=raw_prewait, + n_valid_prewait=self.n_valid_prewait, + ) + + @cute.jit + def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): + num_kv_heads = scores.shape[0] + self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( + grid=(scores.shape[2] * num_kv_heads, 1, 1), + block=(SELECT_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + +# The select reads its row's valid-block count before its grid wait instead of after it, which takes one dependent +# global load off the post-wait path; the selection is unchanged bit for bit. Exact because the count +# (MsaMetadata.msa_n_valid_blocks) is written only before the step's first layer, behind a plain launch +# (attention/test_modeling_v2_m3_n_valid_contract.py). +N_VALID_PREWAIT = True +#: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). +kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} + + +def kernel_key( + use_pdl: bool, early_trigger: bool, trigger_phase: int +) -> Tuple[bool, bool, int, bool]: + """The compile key of a launch (the trigger phase does not matter with the early trigger). The kernel reads its + row's valid-block count before its grid wait (``N_VALID_PREWAIT``).""" + return ( + bool(use_pdl), + bool(early_trigger), + 0 if early_trigger else int(trigger_phase), + N_VALID_PREWAIT, + ) + + +def _compile(key: Tuple[bool, bool, int, bool]): + if key not in kernel_cache: + heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() + scores = make_fake_tensor( + cutlass.Float32, + (heads, blocks, rows), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) + out = make_fake_tensor( + cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) + ) + kernel_cache[key] = cute.compile( + _SelectLongDslKernel(*key), + scores, + n_valid, + out, + cutlass.Int32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + return kernel_cache[key] + + +def select_long_dsl( + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool, + early_trigger: bool, + trigger_phase: int, +) -> None: + """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" + _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( + scores, n_valid_blocks, out, init_blocks, local_blocks + ) + + +def prebuild(keys, device) -> list: + """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" + scores = torch.zeros( + 1, 2048, 4, dtype=torch.float32, device=device + ) # capacity: fixed: a dummy prebuild row, not a size + n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) + out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) + built = [] + for key in keys: + key = kernel_key(*key) + _compile(key)(scores, n_valid, out, 0, 1) + built.append(key) + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_shard_accept.py new file mode 100644 index 000000000000..2b7d05ad920f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/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 reached through the ``ops/spec`` entry ``m3_shard_accept``. +""" + +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from ..ops.spec.m3_shard_accept import m3_shard_accept as _shard_accept_op + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 # capacity: fixed: DECODE_MAX_TOKENS verify rows + + +@dataclass(eq=False) +class AcceptWorkspace: + """Per-rank exchange slots ([2 halves][SHARD_ROWS rows][world] uint64) behind one multicast mapping, the local + key accumulators and the (finished CTAs, call) counters.""" + + uc: torch.Tensor + mc_ptr: int + keys: torch.Tensor + counters: torch.Tensor + rank: int + world_size: int + handle: Any = None + comm: Any = None + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + if not _mnnvl_workspace_all_succeeded(comm, True): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = mc.data_ptr(), mapping.tp_rank + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None + return AcceptWorkspace( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) + + +def _workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace, created collectively outside CUDA-graph capture on first use.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = create_workspace(mapping, device, SHARD_ROWS) + return _WS + + +def shard_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, +): + """Greedy verification from this rank's BF16 logits shard. + + Args: + logits: [nreq * (draft_len + 1) <= SHARD_ROWS, V / tp] BF16, unit column stride: this rank's columns + (global ids ``vocab_offset ..``) of the verify rows, request-major. + draft_tokens: [nreq, draft_len] int32, the drafts being verified. + batch_idx: [nreq] int32, each request's row in the verify layout. + force_total / force_vec: forced acceptance (see ``m3_shard_accept``). + next_new: optional int32 [nreq, >= 1]; column 0 receives the last accepted token. + + Returns: + (accepted_tokens [nreq, draft_len + 1] int32, num_accepted [nreq] int32, first-pass gather ids [nreq] + int64), or None when the shape is not served or the workspace is missing during a capture. Collective over + the TP group. + """ + nreq, draft_len = draft_tokens.shape + if ( + logits.dim() != 2 + or logits.shape[0] != nreq * (draft_len + 1) + or logits.shape[0] > SHARD_ROWS + or logits.dtype != torch.bfloat16 + or logits.stride(1) != 1 + ): + return None + ws = _workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + _shard_accept_op( + logits, + vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_tokens.contiguous(), + batch_idx, + accepted, + num_accepted, + gather_ids, + force_total, + next_new, + _num_sms(), + force_vec, + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_spec_tail.py new file mode 100644 index 000000000000..0d4dc9cd1b64 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_spec_tail.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small fused kernels of the MiniMax-M3 Eagle3 step (GB300, decode). + +``chunk_norm`` / ``capture_norm`` apply the drafter's per-layer fc_norm RMSNorms to the captured target hidden states +(the in-tree ``apply_eagle3_fc`` norms, same fp32 arithmetic); ``draft_input`` is a fused draft pass's input (the +embedding lookup, the drafter layer's input_layernorm / hidden_norm and their concat, same arithmetic); +``step_gather`` / ``slot_scatter`` move the step's integer outputs between the spec sampler's store and the engine's +input buffers. +""" + +from typing import Optional, Tuple + +import torch + +# The kernels are compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu) and reached through the +# ops/spec entries. +from ..ops.spec.m3_capture_norm import m3_capture_norm as _capture_norm_op +from ..ops.spec.m3_chunk_norm import m3_chunk_norm as _chunk_norm_op +from ..ops.spec.m3_draft_input import m3_draft_input as _draft_input_op +from ..ops.spec.m3_slot_scatter import m3_slot_scatter as _slot_scatter_op +from ..ops.spec.m3_step_gather import m3_step_gather as _step_gather_op + + +def draft_input( + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + eps_embed: float, + w_hidden: torch.Tensor, + eps_hidden: float, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(cat([rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden], -1) [M, 2H], hidden[rows] [M, H]) for + int32 ids [M] and int64 rows [M] (rows None: hidden[:M]); ids outside the table embed to zero.""" + m, h = ids.shape[0], table.shape[1] + out = torch.empty((m, 2 * h), dtype=table.dtype, device=table.device) + residual = torch.empty((m, h), dtype=table.dtype, device=table.device) + return _draft_input_op( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, use_pdl + ) + + +def chunk_norm( + src: torch.Tensor, weights: torch.Tensor, eps: float, use_pdl: bool = True +) -> torch.Tensor: + """Per-chunk RMSNorm of src [M, >= C*H] (row stride may exceed C*H) with weights [C, H] -> [M, C*H].""" + c, h = weights.shape + out = torch.empty((src.shape[0], c * h), dtype=src.dtype, device=src.device) + return _chunk_norm_op(src, weights, eps, out, use_pdl) + + +# The overlap scheduler's per-step gathers and the spec sampler's slot stores as one kernel each (step_gather / +# slot_scatter) instead of ~10 and 4 small torch ops; called by the target's SpecStepFold. + + +def step_gather( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 +) -> bool: + """See step_gather_kernel; False when the inputs are outside what the kernel covers (caller falls back).""" + ts = (new_tok, next_draft, lens, slots, ppi, ids, draft, pos_off, kv_off) + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or next_draft.dim() != 2 + or new_tok.dim() < 2 + ): + return False + if (new_tok.dim() == 3 and new_tok.shape[2] != 1) or not ( + lens.is_contiguous() + and ids.is_contiguous() + and draft.is_contiguous() + and pos_off.is_contiguous() + and kv_off.is_contiguous() + ): + return False + _step_gather_op( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 + ) + return True + + +def slot_scatter( + slots, src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft +) -> bool: + """See slot_scatter_kernel (2-D int32 views, zero padding); False to fall back.""" + ts = (src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft) + if any(t.dtype != torch.int32 or not t.is_cuda for t in ts) or slots.dtype != torch.long: + return False + two = (src_new, dst_new, src_next_new, dst_next_new, src_draft, dst_draft) + if any(t.dim() != 2 for t in two) or src_lens.dim() != 1 or dst_lens.dim() != 1: + return False + _slot_scatter_op( + slots, + src_new, + dst_new, + src_next_new, + dst_next_new, + src_lens, + dst_lens, + src_draft, + dst_draft, + ) + return True + + +def capture_norm( + h: torch.Tensor, + residual: Optional[torch.Tensor], + weight: torch.Tensor, + eps: float, + dst: torch.Tensor, +) -> None: + """dst [M, H] (row stride may exceed H) = rmsnorm(h + residual) * weight with chunk_norm's rounding.""" + _capture_norm_op(h, residual, weight, eps, dst) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp4/kernels/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_tokmap_guard.py new file mode 100644 index 000000000000..db0077117fe3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_tokmap_guard.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The token map of ``thop.trtllm_gen_generation_preprocess`` for one request with several tokens. + +For multi-token generation the op's QKV preprocessing places every row by the workspace token map (``tokensInfo``, +(request, token) per row), but builds that map only when ``isBuildDecoderInfoKernelNeeded()`` holds, which at one +request with the default RoPE it does not. The op then applies whatever map the last multi-request call left in the +workspace: a row can land on another request's length and pages, and its RoPE position and K/V slot are wrong. This +module wraps the op so that, for one request and several tokens, the map the build would write, (0, t) for each row t, +is written into the op's own map slice of the workspace (the offset from the op's layout function) just before the +call. Every other call passes through unchanged. The rows' positions still come from the request's generation length, +exactly as with a freshly built map. + +The write is a device copy from a map prebuilt at install (no allocation, no host sync), so it is captured into the +decode graphs with the call it precedes. +""" + +from __future__ import annotations + +from typing import Callable, Dict, Optional, Tuple + +import torch + +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.logger import logger + +#: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). +_ARG_NAMES = ( + "qkv_input", + "workspace", + "sequence_lengths", + "spec_decoding_generation_lengths", + "spec_decoding_position_offsets", + "kv_cache_block_offsets", + "host_kv_cache_pool_pointers", + "host_kv_cache_pool_mapping", + "kv_scale_orig_quant", + "kv_scale_quant_orig", + "attention_output_orig_quant", + "rotary_inv_freq", + "rotary_cos_sin", + "mrope_position_deltas", + "layer_idx", + "seq_offset", + "num_heads", + "num_kv_heads", + "head_size", + "tokens_per_block", + "kv_cache_quant_mode", + "max_attention_window_size", + "cyclic_attention_window_size", + "num_tokens", + "batch_beam", + "input_seq_length", + "max_past_kv_length", + "rotary_embedding_dim", + "rotary_embedding_base", + "rotary_embedding_scale_type", + "rotary_embedding_scale", + "rotary_embedding_max_positions", + "position_embedding_type", + "bmm1_scale", + "bmm2_scale", + "fp8_context_fmha", + "predicted_tokens_per_seq", + "attention_chunk_size", + "multi_processor_count", + "total_num_blocks", + "kv_factor", + "need_build_kv_cache_metadata", + "cross_attention", + "skip_fmha_workspace", +) +_INT2_BYTES = 8 # capacity: fixed: an element size + +_original: Optional[Callable] = None +_max_rows = 0 +_maps: Dict[torch.device, torch.Tensor] = {} +_offsets: Dict[Tuple, int] = {} + + +def _no_taken(rows: int) -> None: + pass + + +def _no_missed(rows: int, reason: str) -> None: + pass + + +_on_taken: Callable[[int], None] = _no_taken +_on_missed: Callable[[int, str], None] = _no_missed + + +def _indexed(device: torch.device) -> torch.device: + """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: + """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" + device = _indexed(device) + m = _maps.get(device) + if m is None or m.shape[0] < rows: + reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" + _on_missed(rows, reason) + raise RuntimeError(f"m3_tokmap_guard: {reason}") + return m[:rows] + + +def _map_offset(args: dict, rows: int) -> int: + x = args["qkv_input"] + skip = bool(args.get("skip_fmha_workspace", False)) + key = ( + x.dtype, + rows, + int(args["num_heads"]), + int(args["head_size"]), + int(args["rotary_embedding_dim"]), + int(args["num_kv_heads"]), + skip, + ) + off = _offsets.get(key) + if off is None: + layout = thop.get_trtllm_gen_generation_workspace_layout( + x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip + ) + off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) + if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: + raise RuntimeError( + f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " + f"(offset {off}, {size} bytes): the layout changed under this guard" + ) + _offsets[key] = off + return off + + +def _guarded(*args, **kwargs): + a = dict(zip(_ARG_NAMES, args)) + a.update(kwargs) + rows = int(a["num_tokens"]) + if ( + a.get("spec_decoding_generation_lengths") is not None + and int(a["predicted_tokens_per_seq"]) > 1 + and int(a["batch_beam"]) == 1 + and rows > 1 + and not bool(a.get("cross_attention", False)) + ): + ws = a["workspace"] + src = _one_request_map(ws.device, rows) + off = _map_offset(a, rows) + # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. + dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] + dst.view(torch.int32).view(rows, 2).copy_(src) + _on_taken(rows) + return _original(*args, **kwargs) + + +def install( + device: torch.device, + max_rows: int, + on_taken: Callable[[int], None], + on_missed: Callable[[int, str], None] = _no_missed, +) -> bool: + """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to + ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this + call installed the wrapper.""" + global _original, _max_rows, _on_taken, _on_missed + device = _indexed(device) + if device not in _maps or _maps[device].shape[0] < max_rows: + t = torch.arange(max_rows, dtype=torch.int32, device=device) + _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() + _max_rows = max(_max_rows, max_rows) + _on_taken = on_taken + _on_missed = on_missed + if _original is not None: + return False + _original = thop.trtllm_gen_generation_preprocess + thop.trtllm_gen_generation_preprocess = _guarded + logger.info( + "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " + f"(up to {max_rows} rows)" + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp4/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/__init__.py new file mode 100644 index 000000000000..ef3a4c5847ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) as a CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..932971278c49 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows per CTA = the MMA's M # capacity: fixed: a kernel tile constant +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row # capacity: fixed: a kernel tile constant +MMA_INST_K = 32 +K_BLOCKS = ( + CTA_K // MMA_INST_K +) # 4 MMAs per k-tile, one per scale byte of the atom # capacity: fixed: a kernel tile constant +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +# The swigluoai activation of MiniMax-M3's dense MLP, gate * sigmoid(1.702 * gate) * (up + 1) with gate clamped to <= 7 +# and up to [-7, 7], in exactly the instruction sequence the Triton silu_and_mul kernel compiles to (modules/swiglu.py +# on SM107): the gate clamp on the packed bf16 pair, the up clamp as min.xorsign.abs, two rounded multiplies into +# ex2.approx, div.full for the reciprocal, two rounded multiplies, round to bf16. No multiply feeds an add, so the +# explicit .rn forms compile to what the Triton kernel's contractable forms do. +SWIGLU_OAI_ALPHA = 1.702 +SWIGLU_OAI_LIMIT = 7.0 + + +@dsl_user_op +def _swiglu_oai_bf16x2(gate_word, up_word, *, loc=None, ip=None): + """The activation of one packed bf16 pair of gate and up values: its two bf16 results as one int32 word (the low + half from the low halves).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [cutlass.Int32(w).ir_value(loc=loc, ip=ip) for w in (gate_word, up_word)], + "{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, $1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, $2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 $0, r1, r0;\n\t}", "=r,r,r", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[ + bool + ], # x_bf16: x_words is [M, 2K] gate | up; the activation is swigluoai(gate, up) +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in if x_swiglu else k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + if cutlass.const_expr(x_swiglu): + # The up half of the same columns, K values after the gate in the row. + qz_wvecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(k_in // 2 + qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + if cutlass.const_expr(x_swiglu): + qz_word = _swiglu_oai_bf16x2( + qz_word, cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + ) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[bool] = False, +) -> None: + # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. + assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( + "x_swiglu needs x_bf16 and excludes din" + ) + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, x_swiglu, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/op.py new file mode 100644 index 000000000000..3e6b07326658 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/op.py @@ -0,0 +1,521 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) for already-quantized activations of at most 8 tokens. + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring): GB300 exposed time behind a 4.25 us PDL predecessor at M 1 / 8, vs the fused low-M GEMM: +# QKV 2.78 / 3.16 vs 4.58 / 6.21 us, o_proj 2.23 / 2.41 vs 4.40 / 5.09, shared gate_up 2.43 / 2.74 vs 3.56 / 4.30, +# shared down 1.74 / 1.66 vs 2.92 / 3.38. A config is (split, ring[, contig, nacc, a_tmem, st_async]); the optional +# fields default to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + (2560, 6144): (4, 12), + (6144, 2048): (2, 8), + # Shared-expert gate_up: split 2, ring 12, single accumulator, A operand from shared memory. + (1536, 6144): (2, 12, 1, 1, 0, 1), + # Shared-expert down: unsplit, ring 3, single accumulator, A operand from shared memory. + (6144, 768): (1, 3, 1, 1, 0, 1), + (2304, 6144): (4, 12), + (6144, 3072): (2, 12), + # Dense-layer gate_up: 6.4 / 6.3 us at M 1 / 8 vs FlashInfer 8.8. + (6144, 6144): (2, 12), +} +# Per-token-count configs (N, K, M): the grid that suits the concurrent routed-expert engine depends on how long +# that engine runs. At 2 tokens the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = { + (1536, 6144, 2): (4, 12, 1, 1, 0, 1), +} +# Above 8 tokens (spec-decode verify rows): the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = { + (1536, 6144): (4, 12, 1, 1, 0, 1), +} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # GB300 exposed time vs the FlashInfer prequantized GEMM (behind a PDL predecessor), at 24 / 48 tokens: QKV 3.76 / + # 4.63 vs 7.46 / 11.44 us, o_proj 2.74 / 3.33 vs 4.46 / 5.84, shared gate_up 4.57 / 5.71 vs 7.16 / 11.25. Behind + # only for K = 768 (unsplit shared down) at 17-32 tokens: 3.1 vs 2.9. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], + x_swiglu: bool = False, +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + if nacc != 1: + # More than one TMEM accumulator is broken on SM107 (nondeterministic or garbage + # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + x_swiglu, + ) + return fn + + +def _views( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight, + weight_scale, + y, + x_bf16: bool, + x_swiglu: bool = False, +) -> tuple: + # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. + k = x.shape[1] // 2 if x_swiglu else x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, + x_swiglu: bool = False, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + if x_swiglu: + k //= 2 + n = weight.shape[0] + if weight.shape[1] != k: + return None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu) + if resolved is None: + return None + key, runtime_ar = resolved + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, + x_swiglu: bool = False, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros( + m, + 2 * k if x_swiglu else k, + dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, + device=dev, + ) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_swiglu( + gate_up: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``m3_mx_gemv_bf16(swiglu(gate_up), ...)`` in one launch: ``gate_up`` bf16 [M <= 8, 2K] (contiguous, gate then + up) is activated in the kernel's quantize prologue with the swigluoai arithmetic of the Triton silu_and_mul kernel + (alpha 1.702, limit 7, the (up + 1) offset; bit for bit), then quantized as the bf16 path does.""" + if ( + gate_up.dtype != torch.bfloat16 + or gate_up.dim() != 2 + or not gate_up.is_contiguous() + or gate_up.data_ptr() % 16 != 0 + or gate_up.shape[1] % 2 != 0 + ): + return None + return _run(gate_up, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar, True) + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode.py new file mode 100644 index 000000000000..d146b9d05d2e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode.py @@ -0,0 +1,1730 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +GB300/SM103 tuning targets GQA group 16, top-k 16 and four-token Eagle3 +verification. Split-K, query pairing and launch tables require four KV heads; +the linear-softmax policy also covers one and two KV heads. Other hardware +and shapes retain the default policies, including supported SM100 NVFP4. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* A packed NVFP4 cache is read in place, dequantized in-register from its E2M1 + nibbles and E4M3 block scales. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* A single top-k chunk writes the final output directly, avoiding scratch and + the otherwise redundant LSE merge. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. + +This is the MiniMax-M3 NVFP4 / sm_107 / tp4 target's copy. On top of the +in-tree decode it routes long-context TP4 steps to the one-launch +``m3_attn_lc`` kernel, short-sequence verify steps of at least +``MQ_MIN_TOKENS`` rows to the multi-query producer, releases PDL dependents +early, and has the merge emit the MXFP8 o_proj hand-off (``m3_prequant.SLOT``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + check_decode_span_shape, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +from . import m3_attn_lc, m3_prequant, m3_report + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +_TL_DTYPES = { + torch.bfloat16: tl.bfloat16, + torch.float16: tl.float16, + torch.float32: tl.float32, +} + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size + +# NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale +# byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. +NVFP4_SF_VEC_SIZE = 16 + + +# Fewest query rows for which the merge writes the MXFP8 o_proj hand-off: every decode step, so the MXFP8 o_proj +# never runs its own quantize kernel after attention. +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and swizzled UE8M0 + scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips its own + quantization. One 128-row buffer per width serves every layer: each o_proj consumes it before the next layer's + merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + +MERGE_MXFP8 = MergeMxfp8() + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + if IDENTITY: + # Slot `it` of the identity selection, bounded by this row's own kv extent above. + blk = it + else: + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, + # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + if IDENTITY: + blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) + blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) + else: + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for cand in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + if IDENTITY: + blk_0 = tl.where(valid_0, idx_0, sentinel) + blk_1 = tl.where(valid_1, idx_1, sentinel) + else: + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the + # pool, reads nothing. + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) +# Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) + + +# Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query +# tokens; below it the per-token producer is as fast. +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover +# Warps per multi-query producer program. +MQ_NUM_WARPS = 8 +# The long-context decode (``_longattn_lc_decode``) in its engagement line (no row range: every step the LC +# applies to takes it). +LC_PATH = "m3_sparse_decode_lc" + + +def _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + identity=False, + identity_step=False, +) -> None: + """minimax_m3_sparse_attn_decode through m3_attn_lc: same table, output and MXFP8 hand-off contract. ``identity`` + takes the short-sequence identity selection from the device lengths (the table is not read); ``identity_step`` + marks a step with no select (the kernel reads q after its grid wait), also when the staged table is read.""" + + use_pdl = get_env_enable_pdl() + total_q, num_heads, head_dim = q.shape + linear = _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=1, + gqa_group_size=num_heads, + max_topk=int(topk_idx.shape[-1]), + decode_query_len=decode_query_len, + ) + mxfp8 = None + m3_prequant.SLOT.clear() + if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(num_heads * head_dim, output.device) + # The early trigger lets the o_proj launch and stream its weights while attention runs. + m3_attn_lc.m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + linear_softmax=linear, + mxfp8=mxfp8, + use_pdl=use_pdl, + early_trigger=True, + # This target's op (sparseDecodeLc4): the merge weights formed once per CTA, and the split-K merge dividing + # once per output (fixed order, no atomics; inside the same band of the float64 reference). + merge_once=True, + prefetch=m3_attn_lc.PREFETCH.take(), + next_kv=m3_attn_lc.NEXT_KV.take(), + identity=identity, + identity_step=identity_step, + one_div=True, + ) + m3_report.taken(ONE_DIV_MERGE_PATH, total_q) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), mxfp8[0][:total_q], mxfp8[1] + ) + + +# Engagement path of the identity-selection producers reading their blocks from the kv lengths (IDENTITY). +DEVICE_IDENTITY_PATH = "sparse_decode_device_identity" +# Engagement path of the LC's one-division split-K merge. +ONE_DIV_MERGE_PATH = "sparse_decode_lc_one_div_merge" + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + q may be FP8, as a fused QK-norm/RoPE producer emits it. The kernel widens + it to output.dtype in-register rather than making the caller materialize a + widened copy, which would cost a standalone kernel on every sparse layer. + The attention math is unchanged either way, since E4M3 widens exactly. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. + + Passing the four NVFP4 scale tensors selects the packed-NVFP4 cache instead: + k_paged/v_paged then hold two E2M1 elements per byte, the block scales are + paged E4M3 bytes in the layouts msa_scatter writes, and the two per-tensor + dequant scales are single-element FP32 device tensors. + + num_topk_chunks overrides the split-K factor and exists for tests: the + merged result must not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) + nvfp4_args = (k_block_scale, v_block_scale, k_global_scale, v_global_scale) + kv_nvfp4 = any(arg is not None for arg in nvfp4_args) + scale_cols = head_dim // NVFP4_SF_VEC_SIZE + if kv_nvfp4: + capability = torch.cuda.get_device_capability(q.device) + if capability not in _SM100F_CAPABILITIES: + raise NotImplementedError( + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; " + f"got SM{capability[0] * 10 + capability[1]}." + ) + _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) + # Two E2M1 elements per byte; Triton must see them unsigned so that the + # high-nibble shift does not sign-extend. + k_paged, v_paged = k_paged.view(torch.uint8), v_paged.view(torch.uint8) + k_block_scale, v_block_scale = ( + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused arguments. + scale_arg = kv_scale if use_scale else output + k_sc_arg, v_sc_arg = (k_block_scale, v_block_scale) if kv_nvfp4 else (output, output) + k_gs_arg, v_gs_arg = (k_global_scale, v_global_scale) if kv_nvfp4 else (output, output) + + widen_q = q.dtype in _FP8_DTYPES + if widen_q and output.dtype not in _TL_DTYPES: + raise ValueError( + f"MiniMax-M3 sparse decode cannot widen FP8 q into {output.dtype}; " + f"supported compute dtypes are {sorted(d.__str__() for d in _TL_DTYPES)}." + ) + # Packed NVFP4 dequantizes to FP16. The production producer supplies FP8 + # q, which widens exactly to FP16, so keep both dot operands there instead + # of converting every K/V tile to BF16. Accumulation remains FP32. + if widen_q: + compute_dtype = tl.float16 if kv_nvfp4 else _TL_DTYPES[output.dtype] + else: + compute_dtype = tl.float32 + + # Long-context TP4 steps: attention, split-K merge and MXFP8 hand-off in one launch. + + if ( + kv_nvfp4 + and widen_q + and num_topk_chunks is None + and m3_attn_lc.lc_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ) + ): + _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + # The identity selection from the device lengths (the staged table is not read). + identity=identity_blocks, + # Every identity step (no select ran): q after the grid wait. + identity_step=identity_blocks, + ) + m3_report.taken(LC_PATH, total_q) + return + + query_group_size = ( + _sm100f_nvfp4_query_group_size( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else 1 + ) + # The stock in-tree decode's split-K chunking at every row count (bit for bit its rows), as the SM107 TP2 target. + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = ( + _sm100f_nvfp4_num_topk_chunks( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else None + ) + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one + # program per (request, block) covers all of them and dequantizes each block once. + use_mq = ( + identity_blocks + and total_q >= MQ_MIN_TOKENS + and kv_nvfp4 + and num_kv_heads == 1 + and decode_query_len > 1 + and head_dim == 128 + and gqa_group_size * decode_query_len <= 128 + ) + if identity_blocks and kv_nvfp4 and widen_q: + # The identity-selection route is modeling.SPARSE_DECODE_PRODUCER's. + m3_report.taken( + "sparse_decode_producer", + total_q, + "multi_query" if use_mq else "per_row", + ) + if use_mq: + num_topk_chunks = max_topk + direct_output = num_topk_chunks == 1 + if direct_output: + # With one chunk, acc_o is already normalized and final. Point the + # producer at the caller's output and compile away the LSE store. + # Keeping a real dummy pointer preserves Triton's pointer contract. + o_arg = output + lse_arg = output + o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) + lse_strides = (0, 0, 0) + else: + # Persistent arena rather than torch.empty, so the partials keep one + # address across CUDA graph replays. FP32 keeps the split-K factor from + # perturbing the merged result. + reserve = torch.cuda.is_current_stream_capturing() + o_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + o_strides = tuple(o_arg.stride()) + lse_strides = tuple(lse_arg.stride()) + + use_pdl = get_env_enable_pdl() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = ( + _sm100f_nvfp4_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + if total_q % query_group_size or decode_query_len % query_group_size: + raise ValueError( + f"query group {query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." + ) + + producer_kernel = ( + _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel + ) + device_identity = identity_blocks + if device_identity and not use_mq: + m3_report.taken(DEVICE_IDENTITY_PATH, total_q) + single_kernel_args = ( + { + # No split-K arrival counters: the separate merge kernel combines the chunks. + "merge": (output, None), + "tail": (output.stride(0), output.stride(1), output.stride(2), num_kv_heads), + # An early release pays off when the dependent is the consumer of the + # final output (it can prefetch), or a merge grid that itself releases + # its dependent at entry: the o_proj then launches during this kernel and + # streams its weights while attention runs. Above 8 query rows the + # concurrent weight stream slows the larger attention grid more than the + # prefetch saves (GB300 TP4: -1.1% at 48 rows). + "options": {"FUSED_MERGE": False, "EARLY_TRIGGER": direct_output or total_q <= 8}, + } + if query_group_size == 1 + else {"merge": (), "tail": (), "options": {}} + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + else {} + ) + if use_mq: + from .triton_sparse_decode_mq import launch_mq_producer + + launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + sm_scale=sm_scale, + decode_query_len=decode_query_len, + use_pdl=use_pdl, + early_trigger=False, + num_warps=MQ_NUM_WARPS, + ) + else: + producer_kernel[ + (total_q // query_group_size) * num_topk_chunks, + num_kv_heads, + ]( + q, + k_paged, + v_paged, + scale_arg, + k_sc_arg, + v_sc_arg, + k_gs_arg, + v_gs_arg, + topk_idx, + o_arg, + lse_arg, + *single_kernel_args["merge"], + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + k_sc_arg.stride(0), + k_sc_arg.stride(1), + v_sc_arg.stride(0), + v_sc_arg.stride(1), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + *o_strides, + *lse_strides, + block_table.stride(0), + block_table.shape[1], + k_paged.shape[0], + *single_kernel_args["tail"], + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + WIDEN_Q=widen_q, + COMPUTE_DTYPE=compute_dtype, + USE_PDL=use_pdl, + KV_NVFP4=kv_nvfp4, + USE_LINEAR_SOFTMAX=use_linear_softmax, + SCALE_COLS=scale_cols, + DIRECT_OUTPUT=direct_output, + # Identity selection from the lengths this kernel reads, not the staged table (dispatch kill switch). + IDENTITY=device_identity, + **single_kernel_args["options"], + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + m3_prequant.SLOT.clear() + if not direct_output: + width = num_heads * head_dim + mxfp8 = ( + MERGE_MXFP8.buffers(width, output.device) + if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS + and head_dim == 128 + and output.dim() == 3 + and output.is_contiguous() + and output.dtype == torch.bfloat16 + else None + ) + mxfp8_args = ( + { + "q8_ptr": mxfp8[0], + "sf_ptr": mxfp8[1], + "num_sf_k_tiles": width // 128, + "MXFP8_OUT": True, + } + if mxfp8 is not None + else {} + ) + merge_launch_options = ( + _sm100f_nvfp4_merge_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_arg, + lse_arg, + output, + head_dim, + *o_strides, + *lse_strides, + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + # Release the o_proj at the merge's entry so it streams its weights during the merge. + EARLY_TRIGGER=True, + **mxfp8_args, + **pdl_launch, + **merge_launch_options, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output.view(total_q, width), mxfp8[0][:total_q], mxfp8[1]) + + +__all__ = [ + "MERGE_MXFP8", + "MERGE_MXFP8_MIN_TOKENS", + "MergeMxfp8", + "NVFP4_SF_VEC_SIZE", + "SPARSE_BLOCK_SIZE", + "_sm100f_nvfp4_merge_launch_options", + "_sm100f_nvfp4_num_topk_chunks", + "_sm100f_nvfp4_query_group_size", + "_sm100f_nvfp4_use_linear_softmax", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py new file mode 100644 index 000000000000..997280727b01 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py @@ -0,0 +1,5941 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_107 / tp4 -- self-contained modeling code. + +MiniMax-M3's language model: 60 layers, hidden 6144, 64 query heads and 4 KV heads of dim 128. Layers 0-2 are dense +attention with a dense SwiGLU-OAI MLP; layers 3-59 are MSA block-sparse attention (an index branch scores 128-token +blocks and the top 16 are attended) with a MoE of 128 routed experts (top 4, sigmoid scores with a correction bias, +renormalized and scaled by 2) plus one shared expert. The checkpoint is a modelopt mixed-precision export: NVFP4 routed +experts, MXFP8 projections, fp32 router, bf16 embedding and head; the KV cache is NVFP4 for the sparse layers and FP8 +for the dense ones. + +``tp4`` is four Rubin GPUs (SM 10.7): head-split attention (16 query heads, one KV head and one index head per rank) and +tensor-parallel experts (every expert on every rank, intermediate 768 per rank), no attention DP. Speculative +decoding, when configured, is one-model Eagle3 (draft length 3): the BF16 drafter (``M3Eagle3Draft``) and the BF16 +target head, both at checkpoint precision; decode steps of up to 8 requests verify and draft through fused kernels +at that precision (``M3Eagle3Worker``). + +**The core** (``MiniMaxM3Tp4ModelingV2``) holds the weights (``weights.MODEL_WEIGHTS``) and the runtime resources the +steps use -- one MSA attention backend per layer, the TP all-reduce communicator, the RoPE table, the MXFP8 GEMM +dispatch state the engine tunes -- and dispatches each step, on the host, from its context count alone +(``_select_target``): + +* ``DecodeTarget`` takes every pure generation step. Up to ``DECODE_MAX_TOKENS`` tokens (up to 8 requests, each one + token or a draft verify of 1 + 3 rows) with 1-D input ids it runs the decode-step kernels and picks each layer + boundary's collective by the step's token bucket (``DECODE_PLANS``); it runs its own generic layer for the first + layer when the fused embedding does not take the step, for every layer whose output Eagle3 captures, and for every + layer of a wider step. +* ``PrefillTarget`` takes every step with context rows: prefill chunks and mixed steps, every layer in its generic + form over the MSA backend. + +The two targets share no body: each owns its layer, attention and MoE compositions and its ``step_args``, the one +projection of the step metadata its forward reads. What both call are the module-level op adapters, each picking the +kernel for one call from shapes and layouts (``_linear``, ``_qk_norm_rope``, ``_sparse_qkv_producer``, the router and +MoE-runner calls, the all-reduce + norm calls), which read no step metadata themselves. + +Every fast path the decode step can take, the rows each covers (``FAST_PATHS``, each paired with the entry whose +``CELLS`` certify it) and what a step that cannot take one does are below: **uncovered** (outside every covered +range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise +``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). + +**Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the +next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the +long-context decode for the o_proj. It publishes it in ``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 dataclasses import dataclass +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + resolved_mode, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionInputType, + PositionalEmbeddingParams, + RopeParams, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, + write_msa_phase_kv, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + _ensure_draft_vocab_size, + _get_requested_draft_moe_backend, + _set_draft_kv_cache_quant_algo, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + TensorParallelMode, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen, is_torch_compiling +from tensorrt_llm._utils import mpi_rank, prefer_pinned +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import weights as _weights +from .kernels import m3_attn_lc, m3_fold_tail, m3_mc_allreduce, m3_prequant, m3_report, m3_spec_tail +from .kernels import m3_draft_lowm as _draft_lowm_kernel +from .kernels import m3_draft_step as m3_draft_step_kernel +from .kernels import m3_drafter_attn as _drafter_attn_kernel +from .kernels import m3_fullctx_decode as _fullctx_decode_kernel +from .kernels import m3_index_score as _index_score_kernel +from .kernels import m3_index_score_t2sc as _index_score_t2sc_kernel +from .kernels import m3_sampled_accept as m3_sampled_accept_kernel +from .kernels import m3_select_long_dsl as _select_long_dsl_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 m3_tokmap_guard as _tokmap_guard +from .kernels import msa_kv_update as _msa_kv_update_kernel +from .kernels.long_select import long_select_blocks +from .kernels.mx_gemv import op as mx_gemv +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_dense_decode import m3_dense_decode +from .ops.attention.m3_dense_qk_norm_rope_fp8 import m3_dense_qk_norm_rope_fp8 +from .ops.attention.m3_draft_attention import m3_draft_attention +from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope +from .ops.attention.m3_fullctx_decode import m3_fullctx_decode +from .ops.attention.m3_index_score import SPLIT_K as INDEX_SCORE_SPLIT_K +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_index_score_t2sc import m3_index_score_t2sc +from .ops.attention.m3_msa_attention import m3_msa_attention +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 +from .ops.attention.m3_qkv_insert import m3_qkv_insert +from .ops.attention.m3_select_blocks import m3_select_blocks +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_select_long_dsl import m3_select_long_dsl +from .ops.attention.m3_sparse_decode import m3_sparse_decode +from .ops.comm.m3_allreduce import m3_allreduce +from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm +from .ops.embedding.m3_embedding import m3_embedding +from .ops.gemm.m3_bf16_linear import m3_bf16_linear +from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu +from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS +from .ops.gemm.m3_draft_lowm import m3_draft_lowm +from .ops.gemm.m3_fp32_linear import m3_fp32_linear +from .ops.gemm.m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + m3_mx_gemv, +) +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.gemm.m3_mxfp8_pinned_gemm import m3_mxfp8_pinned_gemm +from .ops.gemm.m3_mxfp8_pinned_quantize import m3_mxfp8_pinned_quantize +from .ops.gemm.m3_target_head import MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS +from .ops.gemm.m3_target_head import m3_target_head +from .ops.moe.m3_moe16_engine import M3Moe16State, m3_moe16_engine +from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .ops.spec.m3_draft_input import m3_draft_input +from .ops.spec.m3_draft_silu import m3_draft_silu +from .ops.spec.m3_draft_step import MAX_TOKENS as DRAFT_STEP_KERNEL_MAX_TOKENS +from .ops.spec.m3_draft_step import m3_draft_step +from .ops.spec.m3_force_vec import m3_force_vec +from .ops.spec.m3_force_vec_req import m3_force_vec_req +from .ops.spec.m3_shard_accept import m3_shard_accept +from .ops.spec.m3_spec_offsets import m3_spec_offsets +from .ops.spec.m3_step_prologue import m3_step_prologue + +_SM = (10, 7) # capacity: fixed: the device capability, not a size +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The largest generation step the decode target takes: 8 requests verifying 1 + 3 drafts and their padding rows. +DECODE_MAX_TOKENS = 48 # capacity: fixed: the decode kernels' step bound + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 tp4 measurements in the comments).""" + + name: str + max_tokens: int + # The attention boundary's multicast all-reduce + residual + RMSNorm (5.05 vs 5.27 us native at 1 token). + attn_mc_ar: bool + # The same kernel at the MoE boundary with the shared + routed add folded in: wins up to 48 tokens. + moe_mc_ar: bool + # The MoE-boundary all-reduce also emits the next QKV activation in MXFP8 for the decode GEMV / GEMM. + qkv_prequant: bool + + +# tiny: every step up to 8 tokens (c1..c8 without speculation, the c1 verify, every draft pass); small: the c2 +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (kernels/m3_mc_allreduce.BUCKETS). +# SM 10.7 TP4: one bucket. MC_AR_TP4 puts every decode boundary on the multicast all-reduce + residual + RMSNorm +# (trtllm::minimax_m3_sm107_tp4_mc_allreduce_norm, its row kernel: bitwise to the stock TP4 oneshot all-reduce fusion +# on the residual and the normed rows) with the MXFP8 copy for the next QKV; False = the built-in all-reduce. +MC_AR_TP4 = True +# Up to this many rows the decode all-reduces (not the MoE finalize) take the 8-CTA cluster kernel instead of the row +# kernel: the same stock oneshot order (opts bit 32: per-column sum of squares, warp butterflies, 6-warp groups, 4 group +# sums from 0.f), so bitwise the row kernel's output, with each row's exchange spread over 8 CTAs. Rubin TP4 synth, both +# orders: c1 (4 rows) -12, c2 (8 rows) -38 us/step, but c4 / c8 (16 / 32 rows) +5 / +45, where +# 8x the CTAs per row cost more than they save. 0 = the row kernel at every row count. +MC_AR_CLUSTER_MAX_ROWS = 8 # capacity: fixed: measured cluster-kernel crossover +# Every decode / drafter boundary all-reduce launches with the max-smem carveout preference (the cluster kernel's +# option 8 at every row count, the row kernel's carveout instance), so the next kernel's CTAs (the QKV GEMV's staged +# tiles) can co-reside early. Launch attribute only: bitwise unchanged. Kill: False (the GB option table). +MC_AR_CARVEOUT = True +# The fused draft loop's boundary all-reduces (M3Eagle3Worker._m3_draft_ar_norm) launch WITHOUT that carveout: the +# kernel after the drafter's last boundary is the draft step (one CTA per SM, ~145 KB of shared memory) rather than a +# staged GEMV, and TP4 alone of the sm_107 targets ran its draft pass ~74-80 us in-model vs ~57 us standalone at world 4 +# in the step census. Launch attribute only: bitwise unchanged. Kill: MC_AR_CARVEOUT (the shared setting). +DRAFT_MC_AR_CARVEOUT = False +# The fused draft loop's two all-reduce boundaries per pass (o_proj -> post-attention norm, down -> final norm) of up +# to this many rows run the multicast all-reduce + residual + RMSNorm's row kernel in its drafter form (opts bit 16: +# FlashInfer fused_add_rmsnorm's reduction order on the unrounded FP32 sum, so bitwise the drafter's stock plain +# all-reduce + RMSNorm) instead of the built-in all-reduce and the separate norm. Kill: 0. +DRAFT_MC_AR_MAX_ROWS = ( + 8 if MC_AR_TP4 else 0 +) # capacity: fixed: the drafter boundary kernel's certified rows +# Those drafter boundaries run the 8-CTA cluster kernel's drafter form (opts bit 128; TP8's drafter boundary kernel) +# instead of the 1-CTA row kernel. Exact-class: deterministic, the updated residual bitwise, the RMS sum of +# squares in the cluster's order instead of FlashInfer's, so a normed row can differ by one BF16 rounding step (7/64 +# rows at 1 row, 14/256 at 4). Rubin world 4, both orders: 4.39 -> 3.75 us per boundary at 1 row, +# 4.79 -> 4.14 at 4 rows, ~ -2..-3.5 us/step. Kill: False (the row kernel, bitwise the stock boundary). +# OFF: in-model timing was flat and natural acceptance length was grey (lower bound -1.16% vs the -1% line). +DRAFT_MC_AR_CLUSTER = False +# The drafter's o_proj and down projection, on the copied split-K kernel (kernels/m3_lowm_splitk), also push their rows +# into the drafter boundary all-reduce's exchange (the row kernel's slots: the canonical words its own push would store, +# at the (token, 768-column slice) call parity), which then runs prepushed 2 instead of pushing x itself; the copy +# releases its dependents only after its own grid wait and the pushes (the prepushed >= 2 contract). Bitwise. The TP8 +# port: ~ -4 us/step at world 8. Kill: False. +DRAFT_AR_PUSH = True +# The drafter boundary all-reduce fed by its projection's push (DRAFT_AR_PUSH): its own taken / not-taken lines. +DRAFT_AR_PUSH_PATH = "m3_draft_ar_push" +DECODE_PLANS = ( + DecodePlan( + "tiny", DECODE_MAX_TOKENS, attn_mc_ar=MC_AR_TP4, moe_mc_ar=MC_AR_TP4, qkv_prequant=MC_AR_TP4 + ), +) + +# The dense layers' decode attention (trtllm-gen, msa_decode.dense_decode_max_seq_len): the max_seq_len it is handed +# per request count. It only sizes the multi-CTA KV split (every CTA still covers its share of the actual length; the +# result changes only by the split-merge order), and the model maximum over-splits long-context caches at one +# request: GB300 TP4, 16 heads, FP8 KV, 64K context: 15.1 us at 16K vs 28.6 us at 1M. Two requests are neutral (the +# grid is unchanged) and keep the step's bound. +# The generic layer's boundaries (prefill, mixed steps, Eagle3 capture layers): the most rows that take the multicast +# all-reduce + residual + RMSNorm, alone and with the shared + routed add folded in. TP4 skeleton: none (the built-in +# all-reduce) until the multicast all-reduce is certified on Rubin. +GENERIC_MC_AR_MAX_TOKENS = 0 # capacity: fixed: off at TP4 (the built-in all-reduce) +GENERIC_MC_AR_FOLDED_MAX_TOKENS = 0 # capacity: fixed: off at TP4 (the built-in all-reduce) +# An Eagle3 capture layer whose capture norm is armed (CAPTURE_NORM) has its fused feed-forward boundary all-reduce +# (the 8-CTA cluster kernel up to MC_AR_CLUSTER_MAX_ROWS, the row kernel above) also write the capture: +# rmsnorm(updated residual) * fc_norm[i], bit for bit the capture_norm kernel, which then does not run. The MoE +# finalize runs the shared GB300 op and keeps the capture_norm kernel. Wired at the generic layer's fused boundary +# (layers.M3DecoderLayer._feed_forward_boundary) and through M3DecodeTarget._layer's boundaries (``capture``) for a +# capture layer on the fast decode layer. Kill: False (the capture_norm kernel after the boundary). +CAPTURE_NORM_IN_BOUNDARY = True + +# TP4 skeleton: no cap (the built-in split) until a TP4 SM 10.7 cap is measured and certified. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) + +# The dense layers' decode attention on PrimTS split-KV decode (kernels/m3_fullctx_decode) instead of the trtllm-gen +# call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the +# BF16 merge. At TP4 (16 query heads, 1 KV head per rank) the stock split is capped at 16 cluster CTAs per (request, +# KV head), i.e. 16 of 212 CTAs at one request: stock dense fmhaSm107a 33.4-33.8 us per call and the stock drafter +# fmhaSm100f 45.4 us per pass at 165k. Rubin TP4, FP8 KV, 165k, chained us per call +# (microbench), every recipe passing the fp32 oracle: (4, 1) 13.70 (stock 33.6), (1, 1) 11.52. At one +# request the K/V stream is latency-bound, so both deepen the K/V load pipeline (the load pipeline only; the same +# bits): 12 stages at one query row, 8 at four (the long-context sweep: (4, 1) 12.95 / 17.00 / 20.98 us +# per call at 165k / 323k / 512k with 8 stages vs 13.03 / 17.15 / 21.30 with 12; (1, 1) is fastest with 12). Two +# requests (microbench; stock at c2: dense fmhaSm107a 26.5 us, drafter fmhaSm100f 44.0-44.6 us per pass, +# census): (4, 2) 17.34, (1, 2) 14.16. +# K/V stages in flight of the (4, 1) recipe (12 was the previous depth; the same bits at any depth). +FULLCTX_DECODE_Q4_KV_STAGES = 8 # capacity: fixed: a pipeline depth, not a size +FULLCTX_DECODE_RECIPES: Mapping[Tuple[int, int], Mapping[str, Any]] = MappingProxyType( + { + (4, 1): MappingProxyType( + { + "config": { + "tile_size_q": 32, + "groups_tokens_heads_q": True, + "kv_stages": FULLCTX_DECODE_Q4_KV_STAGES, + }, + "splits": 96, + } + ), + (1, 1): MappingProxyType({"config": {"kv_stages": 12}, "splits": 108}), + (4, 2): MappingProxyType( + {"config": {"tile_size_q": 32, "groups_tokens_heads_q": True}, "splits": 48} + ), + (1, 2): MappingProxyType({"config": {}, "splits": 80}), + } +) # capacity: fixed: pinned recipes per certified (rows, requests) +# Kill constant: False puts the dense layers back on the trtllm-gen call at every shape. +FULLCTX_DECODE_ENABLED = True +# The dense layers' FP8 q/k/v (decode steps of the hybrid cache): the op minimax_m3_sm107tp4_dense_qk_norm_rope_fp8 +# (head-count generic) rounds Q/K to BF16 and then to E4M3 in its store, the bits of the BF16 fused_qk_norm_rope output +# cast to FP8. False (or a build without the op) keeps the single-rounding fused_qk_norm_rope FP8 output. +DENSE_QKV_FP8_VIA_BF16 = True +# On a full-context decode step that op also stores the new K/V rows into the layer's paged pool (the bytes +# write_msa_phase_kv's indexed copies store, one launch instead of about ten). False writes them with +# write_msa_phase_kv. +DENSE_KV_INSERT_FUSED = True +# The engagement line of that fused store (taken, one per row count and rank). +DENSE_KV_INSERT_PATH = "m3_dense_kv_insert" +# The drafter MLP's C++ SwiGLU (DRAFT_SILU_CPP, M3Eagle3Draft.mlp): its own taken / not-taken lines. +DRAFT_SILU_PATH = "m3_draft_silu" +# The split-KV merge on the BF16 op (minimax_m3_sm107tp4_fullctx_reduce, head-count generic) instead of PrimTS's FP16 +# reducer and a cast. +FULLCTX_DECODE_OWN_REDUCER = True +# The merge also writes the MXFP8 copy of its BF16 rows (E4M3 + 128x4-swizzled UE8M0, bit for bit +# trtllm::mxfp8_quantize) into the MERGE_MXFP8 buffers and publishes it to m3_prequant.SLOT, so the dense layer's +# MXFP8 o_proj skips its own quantization. Kill: False leaves the o_proj to quantize its BF16 input. +FULLCTX_DECODE_MXFP8 = True +FULLCTX_DECODE_MXFP8_PATH = "m3_fullctx_merge_mxfp8" +# The K/V length bound the full-context plans are compiled for; a step whose cache length bound is longer is a miss. It +# covers the long-context serving config's max_seq_len (1 << 20) plus the speculative rows the cache reserves past it. +FULLCTX_DECODE_MAX_KV_LEN = ( + 1 << 21 +) # capacity: fixed: the plans' KV-length bound (1M context + spec rows) +# The Eagle3 drafter layer's generation phase on the same full-context decode (kernels/m3_drafter_attn: the in-tree +# RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its +# split at 8 CTAs per (request, KV head) on SM 10.7: 45.4 us per draft pass at 165k. Kill constant. +DRAFTER_FULLCTX_ENABLED = True +# On those passes the drafter's RoPE, FP8 query and FP8 K/V append run as one launch (the op +# minimax_m3_sm107tp4_drafter_rope_fp8_kv: the in-tree preprocess's GPT-NeoX rotation and unit-scale E4M3 store, and +# the torch E4M3 cast of its BF16 query) instead of thop.trtllm_gen_generation_preprocess and the cast. Kill constant. +DRAFTER_ROPE_FUSED = True +# The engagement line of that launch (taken, one per row count and rank). +DRAFTER_ROPE_PATH = "m3_drafter_rope" +# thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request +# multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by +# that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, +# token-0 drafts). kernels/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# calls: what the build writes (padded rows included), so bitwise the op with a fresh map. Kill: False. +PREPROCESS_TOKMAP_FIX = True +PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" +# How the rotation's products are summed, as the in-tree preprocess build sums them (0: two products then the add, +# 1: fma(cos, x, sin * pair), 2: fma(sin, pair, cos * x)); the op's bitwise test pins it. +DRAFTER_ROPE_MODE = 1 # capacity: fixed: a launch mode, not a size + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the index scorer launches and stages its pre-wait loads while the producer runs (its +# outputs are read only after the scorer's own grid wait). +PRODUCER_EARLY_TRIGGER = True +# Steps of up to this many rows keep the producer's trigger at its exit instead. On Rubin TP4 the early-launched scorer +# (its grid doubles with the requests, its pre-wait index-K prefetch running) is resident through the producer's whole +# post-QKV segment and slows the producer's dependent loads; from 16 rows the overlap pays. In-model, locked memory +# clock, tactic- and FlashInfer-pinned, A and B in one job: -75 / -176 us per decode step at 1 / 2 requests, +# +67 at 4 with the trigger at the exit. Timing only: the scorer reads the producer's outputs after its own grid wait. +# Kill: 0 (the early trigger at every row count). Set to 0 on the shipped stack: with the long-row select's late release +# and post-load trigger, the scorer's pre-wait overlap with the producer is worth more than the producer's own +# segment (both orders, double-pinned: -12 / -42 us at 1 / 2 requests with 0). +PRODUCER_LATE_TRIGGER_MAX_ROWS = 0 # capacity: fixed: measured trigger crossover (0: off) +# The producer reads its per-step inputs (position ids, cache slots, RoPE rows) before its grid wait, leaving only the +# packed QKV row after it: this target's copy of the producer (trtllm::minimax_m3_sm107_tp4_producer_pw, +# minimaxM3Sm107Tp4/producerPw4.cu, the TP2 producerPw2 at 16 query heads). Their writers are host-to-device copies +# before the replay and plain-launched kernels earlier in the step (the input preprocessing, msa_kv_update or its torch +# fallback), so they are final before the producer can start; a PDL-launched writer of any of them would have to turn +# this off. The call site's packed.contiguous() and position-id copies are plain launches too. Bit-identical outputs. +# Kill: False (the shared minimax_m3_producer_pw at 16 query heads, which reads them after its wait). +PRODUCER_PREWAIT_STEP_INPUTS = True + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. +EMBED_FIRSTNORM_MAX_TOKENS = 8 # capacity: fixed: built rows; wider steps embed generically +# The o_proj pushes its rows into the attention boundary's multicast all-reduce (prepushed mode 2) up to this many +# tokens; wider steps run the GEMV and a plain multicast all-reduce. +ATTN_PUSH_MAX_TOKENS = 8 if MC_AR_TP4 else 0 # capacity: fixed: measured push crossover +# The dense layers' MLP down projection pushes its rows into the post-MLP boundary's multicast all-reduce (prepushed +# mode 2), as the o_proj does into the attention boundary (TP2: those hops end 5.5-5.8 us after the down +# GEMV unpushed vs 2.4-2.5 us pushed). Bitwise equal to the unpushed boundary (the same bf16 words). Kill: False. +DENSE_PUSH_TP4 = True +DENSE_PUSH_MAX_TOKENS = ( + 8 if (MC_AR_TP4 and DENSE_PUSH_TP4) else 0 +) # capacity: fixed: the GEMV push rows +AR_PUSH_MODE = 2 # capacity: fixed: an exchange mode, not a size +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch). It pushes its sum into the +# MoE boundary's multicast all-reduce (prepushed mode 2) on the rows MOE_PUSH covers (with MC_AR_TP4). +MOE_ENGINE_MAX_TOKENS = 8 # capacity: fixed: the engine's built rows +MOE_PUSH_MAX_TOKENS = ( + MOE_ENGINE_MAX_TOKENS if MC_AR_TP4 else 0 +) # capacity: fixed: the engine's rows +# The 9..13-row engine (router, routed and shared experts in one kernel); its bf16 sum goes to the boundary. +MOE16_MIN_TOKENS = 9 # capacity: fixed: the 9..13-row engine's lower bound +MOE16_MAX_TOKENS = 13 # capacity: fixed: the 9..13-row engine's built rows +# From this many tokens the MoE runs unfinalized through the TRTLLM-Gen runner and its finalize (top-k gather, +# routing weights, shared expert add) runs inside the boundary all-reduce. 9..13 take it only when their engine is +# missing (outside require mode). TP4 skeleton: no rows (the built-in MoE finalize and all-reduce). +MC_FINALIZE_MIN_TOKENS = 9 # capacity: fixed: the rows above the engines +MC_FINALIZE_MAX_TOKENS = 0 # capacity: fixed: off at TP4 +# A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel +# over the target lm_head's BF16 shard logits (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 sampled (temperature + top-p, advanced_sampling_mode=no_topk) Eagle3 verify step of up to this many requests +# verifies with the fused sampled acceptance (kernels/m3_sampled_accept) over this rank's logits shard: every request +# count a decode step holds. +SAMPLED_VERIFY_MAX_REQUESTS = SHARD_ACCEPT_MAX_REQUESTS # capacity: fixed: a decode step's requests +# The sampled verify runs the sm_107 tp4 copy of the TP2 sampled-accept kernel +# (trtllm::minimax_m3_sm107tp4_sampled_accept: a row's CTAs launch as one thread-block cluster, bounded record polls). +# Same records and tokens. Kill: False (the shared kernel, trtllm::minimax_m3_sampled_accept: every CTA of the grid +# co-resident, unbounded polls). +SAMPLED_ACCEPT_SM107TP4 = True +# The sm_107 tp4 sampled verify loads each row's sampling parameters (temperature, top-p, Philox seed and offset: step +# inputs, not the head's output) before its PDL dependency wait, overlapping the head GEMM's tail; the same values and +# tokens. Kill: False (they load after the wait; the shared kernel always loads them after). Rubin TP2 microbench: +# verify end 20.48 -> 19.97 us at top-p 0.95. +SAMPLED_ACCEPT_PREWAIT_PARAMS = True +# The sm_107 tp4 sampled verify releases its dependents (the first draft pass's input) right after its dependency wait +# instead of at the acceptance epilogue: the dependent launches early and parks at its own griddepcontrol.wait, which +# returns only once the verify has completed, so the values are the same; it saves the launch latency after the verify +# (TP2: draft_input started 2.0-2.4 us after the verify / fc join). Kill: False. Off until its sampled-column +# in-model pairs read. +SAMPLED_ACCEPT_EARLY_TRIGGER = False +# SA-1R: the sm_107 tp4 sampled verify also multicasts each CTA's largest logits with its round-0 record (82: each +# thread's top 2 pooled, the CTA's top 8 plus the bound of everything it does not list; 9: an exact top-9 list), so a +# row whose first candidate is provably in the nucleus decides in round 0 instead of exchanging a second round. The +# same tokens, accepted counts and Philox draws as the two-round kernel, bit for bit (a fp32 margin m * 2^-21 * S +# covers the summation order; any doubt runs round 1). The prototype was exact on 39,968 rows; real +# long-context serving logits decide all 4 rows in round 0 on 79.5% (82) / 80.8% (9) of steps. Off (0, the two-round +# kernel): exact, but the untimed decode step reads +47 / +65 us slower with it on at TP4 (the earlier dependent +# launch it enables costs more than the saved round). +SAMPLED_ACCEPT_1R_K = 0 # capacity: fixed: a kernel variant, not a size +# The sm_107 tp4 sampled verify's owner walk (the one thread that finds its row's drawn token by a serial +# prefix over its strided tokens) takes those tokens 8 at a time: their weights are computed independently, then +# summed in the same order with the same early exit, so the draw, tokens and Philox draws are the one-token walk's, +# bit for bit. Rubin TP4 bench (one request, top_p 0.95): -1.06 / -1.07 us per call at SA-1R K 0 / 82. Kill: False +# (the one-token walk, the stock instantiation; the shared kernel has no 8-token walk). +# Off until its in-model TP4 pair passes. +SAMPLED_ACCEPT_ILP_WALK = False +# A step verified by the in-tree path takes the fused first drafter inputs (the capture norms and the first pass's +# gather) up to this many requests: every request count a decode step holds. +FIRST_DRAFTER_INPUTS_MAX_REQUESTS = SHARD_ACCEPT_MAX_REQUESTS # capacity: fixed: step requests +# A fused-verify step of up to this many verify rows (one or two requests at draft length 3) computes the target +# logits shard from the 12-bit coded copy of the BF16 head (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. +# Off at TP4 (as at TP2): at the TP4 shard (50,016 x 6144) the in-tree BF16 cuBLAS head is faster than the coded kernel, +# 47.9 vs 56.8 us at 4 and 8 rows (TP2: 93 vs 109). True re-enables the coded head. +TARGET_HEAD_CODED = False +TARGET_HEAD_MAX_ROWS = ( + TARGET_HEAD_KERNEL_MAX_ROWS if TARGET_HEAD_CODED else 0 +) # capacity: fixed: the coded head's built rows (off at TP4) +# On a fused-verify step of up to this many verify rows the drafter's first fc is forked onto the aux stream right +# after the last captured target layer, overlapping the remaining target layers and the head (M3Eagle3Worker). +FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover (0: off) +# A sampled fused verify of up to this many verify rows on the in-tree head GEMM forks the drafter's first fc after the +# head instead of beside it (M3Eagle3Worker), as the coded head's verify already does: the fc's split-K grid otherwise +# holds the SMs first and the head GEMM starts late (TP2: 17.4 us late). Kill: 0. +FC_AFTER_HEAD_MAX_ROWS = 8 # capacity: fixed: measured fork-after-head crossover +# A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (M3Eagle3Worker): per pass the +# draft input in one launch (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 = False +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +# The Eagle3 capture layers (decoder layers whose output the drafter reads) run the fast decode layer, with its +# o_proj / MoE / dense-down pushes into the boundary all-reduces, and capture the MoE / dense boundary's reduced +# residual sum (the same BF16 words the generic layer's fused boundary captures). Kill: False (the generic layer, +# no boundary pushes). Rubin TP2: the capture layers' hops end 4.4-5.3 us after their producing +# GEMV unpushed vs 2.4-2.5 us pushed. +CAPTURE_LAYERS_FAST = True +# An Eagle3 capture layer (one with a next norm) ends at the fused boundary (all-reduce + residual + next RMSNorm in +# one kernel) and captures that kernel's reduced residual sum. Kill: False keeps the feed-forward module's own +# all-reduce, captures hidden + residual and runs the next norm separately, as the built-in model does; it also +# turns CAPTURE_LAYERS_FAST off, which captures the fused boundary's sum. +CAPTURE_FUSED_BOUNDARY = True +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) run FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm, the TP2 target's port with the TP4 table) at the rows its table is certified for, up to this +# many rows; other calls keep the in-tree cuBLAS GEMM (split-K + reduce). Kill: 0. Rubin TP4 sweep (locked memory +# clock) at 1 row: qkv 8.15 -> 5.90, o 5.36 -> 3.40, gate_up 12.93 -> 9.93, down 7.68 -> 6.16 us. +DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the low-M tactic table's certified rows +# The drafter MLP's activation (the plain SwiGLU of its in-tree GatedMLP) runs this target's C++ op +# (ops/spec/m3_draft_silu, minimaxM3Sm107Tp4/siluAndMul.cu) instead of the Triton silu_and_mul: bit for bit, one launch +# either way (no Triton on the generation decode path). The fused draft loop's MLP (M3Eagle3Draft.mlp) routes the +# call; the in-tree draft path keeps the Triton op. Kill: False. +DRAFT_SILU_CPP = True +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(m for m in rows if m <= DRAFT_LOWM_MAX_ROWS) + for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() +} +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (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 = 0 # capacity: fixed: off at TP4 (the DSL select takes the long rows) +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) # capacity: fixed: compiled keys per thread +# The long-row select releases the sparse decode after it right after its own grid wait (the decode stages q and the +# request length before its wait) up to this many query rows. On Rubin TP4 the early-launched attention CTAs crowd the +# select and the o_proj GEMV already at 2 requests (8 rows), so the early release is off at every row count: in-model, +# locked memory clock, pinned tactics, long-context draws, 2 instances vs the stored baseline, the replay period with +# the trigger phase below is -209 us per step at 2 requests (mean of draws), -39 at 4 and -12 at 8. +SELECT_EARLY_MAX_ROWS = 0 # capacity: fixed: measured early-release crossover (0: off) +# The long-context sparse decode runs this target's copy of the op (trtllm::minimax_m3_sm107_tp4_sparse_decode_lc, +# minimaxM3Sm107Tp4/sparseDecodeLc4.cu): its split-K merge forms the 16 weights once per CTA (16 threads, through shared +# memory) instead of in each of the 128 threads. Same values in the same order: bitwise the image op's output. The +# SM107 TP2 sparse decode's -0.17 / -0.23 us per layer at 1 / 2 requests. Kill: False (the image op). +SPARSE_DECODE_LC_MERGE_ONCE = True +# With SPARSE_DECODE_LC_MERGE_ONCE, that op's CTAs also pull the layer's o_proj weight and scales into L2 before +# their grid dependency (cp.async.bulk.prefetch.L2, a hint: no value changes), so the o_proj GEMV after it finds them +# L2-warm. The o_proj's own pre-wait staging streams part of them already; the cold-vs-hot GEMV gap left is ~0.18 us +# per layer. Kill: False. +SPARSE_DECODE_LC_PREFETCH_OPROJ = True +# With SPARSE_DECODE_LC_MERGE_ONCE, that op also records each step's selected block ids per layer and pulls into L2 +# the K/V pages + scales of the blocks the NEXT sparse layer selected at the previous step (65-68% of a layer's +# distinct selected blocks recur one step later, measured on all 4 ranks), so that layer's sparse decode reads +# them L2-warm. Hints only; up to N1_MAX_ROWS query rows. Kill: False. +SPARSE_DECODE_LC_PREFETCH_NEXT_KV = True +N1_MAX_ROWS = 128 # capacity: fixed: the next-KV id buffer's rows +# With SPARSE_DECODE_LC_MERGE_ONCE, short-sequence identity steps (every request within 16 blocks, +# MSA_SHORT_SEQUENCE_IDENTITY) also run that op instead of the Triton producer: CTA r attends block r of its request +# when r < min(16, blocks), taken from the device sequence length (the table is not read, so a host-staged identity +# table that misses the in-flight accepted tokens cannot drop a block). Exact-class like every LC row (band vs the +# fp64 reference); removes the Triton decode from the GEN path at TP4. Kill: False (the Triton identity producer). +# Identity steps skip the select, so the producer's entry trigger used to release LC4's pre-wait q read before the +# producer wrote q (NUL tokens in GSM8K outputs). LC4 reads q after the grid wait, which clears that race (0 NUL). +# Matched ON-vs-OFF GSM8K pairs, greedy and sampled, show no accuracy change; acceptance length +0.024. +SPARSE_DECODE_LC_IDENTITY = True +# Up to SELECT_LONG_TRIGGER_MAX_ROWS query rows and without the early release, the long-row select releases the sparse +# decode once every score of its row is loaded (SELECT_LONG_TRIGGER_PHASE 1, before the exact rank of the candidates) +# instead of at its exit (0, the kill value); wider steps keep the exit. The release still follows the select's own +# grid wait, which the sparse decode's pre-wait q / length / block-table reads rely on. +SELECT_LONG_TRIGGER_PHASE = 1 # capacity: fixed: a trigger phase, not a size +SELECT_LONG_TRIGGER_MAX_ROWS = 8 # capacity: fixed: measured trigger crossover +# The long-row select ranks a row's candidates on warp 0 when they number at most 64 (a 64-key bitonic network; the +# other warps exit) and reads the row's valid count before its grid wait; the picks are the same. Kill: False (the +# all-thread rank). In-model, locked memory clock, double-pinned, A and B in one job, both orders: flat within +# noise (mean of draws -5 / -60 us at 1 / 2 requests; replay period at the 323K request +14 / +54); a FLAT-but-exact +# stack item (post-wait select time 1.76 -> 1.54 us per layer at 1 request, 165K tokens). OFF by default: at TP2 the +# warp-0 tail runs ~0.95 us/layer slower in-model next to the early-launched sparse decode (+53 us/step); +# on at TP4 only after an in-model A,B shows a gain. +SELECT_LONG_SL3 = False +# The long-row select runs as kernels/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per row +# instead of one 1024-thread CTA, the same selection bit for bit, with 4, 11 or 24 keys per thread +# (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices (SELECT_LONG_SL3 does not apply). TP4 has the TP8 target's +# per-rank select geometry (1 KV head, 4 rows per request), where on the Rubin decode chain (57 layers behind the +# producer stream, 1 / 2 requests) the select is -0.15 / -0.12 us per layer at 165K tokens, 0.00 / -0.01 at 323K and +# -0.04 / -0.04 at 500K against the 1024-thread kernel. Kill: False (the 1024-thread kernel). +SELECT_LONG_DSL = True +LONG_SELECT_DSL_KEY_WIDTHS = (4, 11, 24) # capacity: fixed: compiled keys per thread +# The MSA index scorer (ops/attention/m3_index_score) prefetches the older index-K pages into L2 before its grid wait +# at every batch; up to this many requests it also issues Q and the first K tile right after the wait and caps that +# prefetch at one block per CTA (larger batches are bandwidth bound, where the extra early tiles only crowd the +# prefetch). On Rubin TP4 the 2-request step keeps the early schedule: with the producer's trigger at its exit +# (PRODUCER_LATE_TRIGGER_MAX_ROWS) the scorer no longer overlaps the producer, and the capped prefetch shortens its +# exposed pre-wait stream. In-model, locked memory clock, tactic- and FlashInfer-pinned, A and B in one job, both +# orders: -260 us per decode step at 2 requests (replay period at the 323K P90 request; -82 mean of draws), flat at +# 1 / 4 / 8. Timing only: the tiles are read after the wait. Kill: 1. +INDEX_SCORE_EARLY_MAX_REQUESTS = 2 # capacity: fixed: measured early-issue crossover +# T2SC: up to 2 requests the sparse decode scores with this target's copy of the sm_107 TP2 / TP8 scorer +# (kernels/m3_index_score_t2sc; ops/attention/m3_index_score_t2sc), whose split is chosen per request count and which +# TMA-loads an all-old first block before its grid wait; the scores are bit-identical +# to the in-tree scorer. TP4 has the TP8 target's per-rank indexer geometry (1 index head over the replicated index-K +# of its KV group), where this schedule is ~100 us per decode step faster at 1 request than split 512 without the +# prestage (TP8 in-model, greedy early column, both orders). Wider steps keep the sm_103-derived scorer above. +# Kill: False (the sm_103-derived scorer at every request count). +INDEX_SCORE_T2SC = True +# At 1 request the split is one wave of CTAs (1272 = 212 SMs x 6 resident at this CTA's shared memory): at the +# 1-index-head geometry (57-layer chain behind the producer, scores and the selection bitwise) -0.03 / -0.27 / +0.06 us +# per layer at 165K / 323K / 500K tokens against 1536. Kill: 1536. +INDEX_SCORE_T2SC_SPLIT_K: Mapping[int, int] = MappingProxyType( + {1: 1272, 2: 768} +) # capacity: fixed: measured splits per request count + +# The decode MXFP8 GEMVs, per projection (N, K): the rows each serves. On an activation its producer quantized +# (the MoE boundary for the QKV, the attention merge for the o_proj) the GEMV serves every row up to 48 but the +# unsplit K=768 shared down projection at 17..32 rows, where FlashInfer's GEMM is faster. +# The QKV MXFP8 GEMVs' (split, ring) at decode rows up to GEMV_QKV_CFG_MAX_ROWS: a TP4-owned override of the shared +# GB300-tuned table (kernels/mx_gemv/op.py). SM107 TP4 sweep +# (4 rows): sparse-layer QKV (8, 8) 2.70 us vs (4, 12) 3.01, dense-attention QKV (8, 8) 2.60 vs 2.83. Split 8 +# sums K in a different fp32 split-K order than split 4: exact-class, not bitwise to the shared config. Kill: False. +GEMV_QKV_SPLIT8_TP4 = True +GEMV_QKV_CFG_MAX_ROWS = 8 # capacity: fixed: the measured override's rows +MX_GEMV_CFG: Dict[Tuple[int, int], tuple] = { + QKV_SPARSE: (8, 8), + QKV: (8, 8), +} # capacity: fixed: (split, ring) per projection +# (exact-class, deterministic) The o_proj GEMV at decode rows up to GEMV_QKV_CFG_MAX_ROWS with strided k-tiles per split +# rank (contig 0: rank r owns k-tiles r, r + split, ...) instead of contiguous ranges, the shared (2, 8): another +# fixed fp32 order. SM107 TP4 4 rows: 2.513 / 2.533 vs 2.550 / 2.563 us; deterministic (3 runs bitwise) and +# max error vs the fp64 product equal to the contiguous order's. Kill: False. +GEMV_OPROJ_STRIDED_TP4 = True +MX_GEMV_CFG_STRIDED: Dict[Tuple[int, int], tuple] = { + O_PROJ: (2, 8, 0) +} # capacity: fixed: (split, ring, contig) per projection + + +def mx_gemv_cfg(n: int, k: int, rows: int) -> Optional[tuple]: + """The TP4 GEMV config override for an [rows, k] x [n, k]^T MXFP8 GEMV call, or None for the shared table's.""" + if rows > GEMV_QKV_CFG_MAX_ROWS: + return None + if GEMV_QKV_SPLIT8_TP4 and (n, k) in MX_GEMV_CFG: + return MX_GEMV_CFG[(n, k)] + if GEMV_OPROJ_STRIDED_TP4: + return MX_GEMV_CFG_STRIDED.get((n, k)) + return None + + +# Synth only (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS, fractional): the fused verify's forced accepted counts are +# drawn per request, keyed by the request's KV length at the verify (its prompt plus what it has accepted), instead of +# from the server-wide counter, so a request's forced counts do not depend on what the server ran before (A/B draws +# become history-independent). Production sampling never forces counts. Kill: False (the shared counter, the in-tree +# sequence). +FORCE_AL_PER_REQUEST = True + +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) # capacity: decode rows +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + QKV_SPARSE: _ALL_ROWS, + QKV: _ALL_ROWS, + O_PROJ: _ALL_ROWS, + SHARED_GATE_UP: _ALL_ROWS, + SHARED_DOWN: _ALL_ROWS - frozenset(range(17, 33)), + DENSE_GATE_UP: _ALL_ROWS, + DENSE_DOWN: _ALL_ROWS, +} # capacity: decode rows per projection +# On a bf16 activation the GEMV quantizes it in its prologue, which pays only up to 8 32-element blocks per thread: +# the most rows each projection takes. +MX_GEMV_BF16_MAX_TOKENS: Dict[Tuple[int, int], int] = { + QKV_SPARSE: 16, + QKV: 16, + O_PROJ: 32, + SHARED_GATE_UP: 16, + SHARED_DOWN: 16, + DENSE_GATE_UP: 8, + DENSE_DOWN: 16, +} # capacity: fixed: measured bf16-GEMV crossover per projection +# MX-TAC (the TP2 target's, ported): the decode-graph MXFP8 projections the GEMVs do not serve (SHARED_DOWN at 17..32 +# rows, rows over DECODE_MAX_TOKENS or a bf16 activation over MX_GEMV_BF16_MAX_TOKENS, or a GEMV decline) take one fixed +# decision instead of the AutoTuner's per-start, per-rank choice: FlashInfer's CuTe DSL quantize and CuTe DSL GEMM. All +# of stock's candidates (quantize {TRT-LLM, CuTe DSL} x GEMM {CUTLASS, CuTe DSL}) give the same bits at the TP4 +# projections (gemm/test_modeling_v2_m3_pinned_gemm.py), so the pin is stock's output on every rank and start. Kill +# constant: False returns these projections to the AutoTuner. +MX_TAC_TP4 = True +# The tactics below are MXFP8QuantizeRunner.CUTE_DSL / FlashInferMXFP8GemmRunner.CUTE_DSL (0). +MXFP8_DECODE_QUANT_TACTIC = 0 # capacity: fixed: a tactic id, not a size +MXFP8_DECODE_GEMM_TACTIC = 0 # capacity: fixed: a tactic id, not a size + +Key = Hashable + + +@dataclass(frozen=True) +class RowRange: + """The shapes one fast path takes, and the catalog entry whose cells certify each of them. + + ``keys`` are the covered shapes: row counts for most paths, (query rows per request, requests) for the attention + kernels, a variant name for thresholds. ``certifies`` maps a cell's spec of ``entry`` to the keys that cell + drives (empty for a cell of another variant). + """ + + path: str + keys: FrozenSet[Key] + entry: OpWrapper + certifies: Callable[[Mapping[str, Any]], Iterable[Key]] + + def covers(self, key: Key) -> bool: + return key in self.keys + + def certified(self) -> FrozenSet[Key]: + return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + + +def _rows(lo: int, hi: int) -> FrozenSet[int]: + return frozenset(range(lo, hi + 1)) + + +def _spec_rows(field: str, **match: Any) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + """A cell's ``spec[field]`` when the cell's other fields match ``match``.""" + + def certifies(spec: Mapping[str, Any]) -> Iterable[Key]: + if any(spec.get(name) != value for name, value in match.items()): + return () + return (spec[field],) + + return certifies + + +def _const(*keys: Key) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + return lambda spec: keys + + +def _sparse_decode_producer(spec: Mapping[str, Any]) -> Iterable[Key]: + if not spec["identity"]: + return () + rows = spec["decode_query_len"] * len(spec["seq_lens"]) + return ("multi_query" if rows >= SPARSE_DECODE_MQ_MIN_TOKENS else "per_row",) + + +def _long_select_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The key widths a cell's rows reach: every width up to its blocks' when rows of mixed lengths are drawn.""" + blocks = spec["num_blocks"] + widest = 2 if blocks <= 2048 else 4 if blocks <= 4096 else 9 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_KEY_WIDTHS if w <= widest) + return (widest,) + + +EMBED_FIRSTNORM = RowRange( + "embed_firstnorm", + _rows(1, EMBED_FIRSTNORM_MAX_TOKENS), + m3_embedding_firstnorm, + _spec_rows("rows"), +) +SPARSE_DECODE_PRODUCER = RowRange( + "sparse_decode_producer", + frozenset({"per_row", "multi_query"}), + m3_sparse_decode, + _sparse_decode_producer, +) +# Every token count of a sparse layer on an NVFP4 KV cache at the TP4 heads (prefill and decode): the producer takes the +# geometry, not a row range. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp4"}), m3_producer, _const("nvfp4_tp4")) +INDEX_SCORE = RowRange( + "m3_index_score", + frozenset({"early", "prefetch"}), + m3_index_score, + lambda spec: ("early" if spec["early"] else "prefetch",), +) + + +def index_score_t2sc_schedule(requests: int) -> Tuple[int, bool]: + """T2SC's (CTAs per request, early issue) at this many requests (a request count in INDEX_SCORE_T2SC_SPLIT_K).""" + return INDEX_SCORE_T2SC_SPLIT_K[requests], requests <= INDEX_SCORE_EARLY_MAX_REQUESTS + + +def _index_score_t2sc_requests(spec: Mapping[str, Any]) -> Iterable[Key]: + """A cell certifies its request count when it drives this target's 1 index head with that count's schedule.""" + requests = len(spec["seq_lens"]) + if spec.get("heads") != 1 or requests not in INDEX_SCORE_T2SC_SPLIT_K: + return () + if (spec["split_k"], spec["early"]) != index_score_t2sc_schedule(requests): + return () + return (requests,) + + +INDEX_SCORE_T2SC_PATH = RowRange( + "m3_index_score_t2sc", + frozenset(INDEX_SCORE_T2SC_SPLIT_K) if INDEX_SCORE_T2SC else frozenset(), + m3_index_score_t2sc, + _index_score_t2sc_requests, +) +LONG_SELECT = RowRange("long_select", frozenset(), m3_select_long, _long_select_widths) + + +def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The 384-thread select's key widths a cell's rows reach (every width up to its blocks' for mixed rows).""" + blocks = spec["num_blocks"] + widest = 4 if blocks <= 4 * 384 else 11 if blocks <= 11 * 384 else 24 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_DSL_KEY_WIDTHS if w <= widest) + return (widest,) + + +LONG_SELECT_DSL = RowRange( + "long_select_dsl", + frozenset(LONG_SELECT_DSL_KEY_WIDTHS) if SELECT_LONG_DSL else frozenset(), + m3_select_long_dsl, + _long_select_dsl_widths, +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS) if MC_AR_TP4 else frozenset(), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=0, drafter=None), +) +# Rows of a fused-draft-loop all-reduce boundary (drafter hidden 6144) on the multicast all-reduce's drafter form. +DRAFT_MC_AR = RowRange( + "m3_draft_mc_allreduce_norm", + _rows(1, DRAFT_MC_AR_MAX_ROWS), + m3_mc_allreduce_norm, + _spec_rows("m", drafter=True), +) +DENSE_PUSH = RowRange( + "dense_push", + _rows(1, DENSE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +ATTN_PUSH = RowRange( + "attn_push", + _rows(1, ATTN_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MOE_PUSH = RowRange( + "moe_push", + _rows(1, MOE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MC_FINALIZE = RowRange( + "mc_finalize_allreduce_norm", + _rows(MC_FINALIZE_MIN_TOKENS, MC_FINALIZE_MAX_TOKENS), + m3_mc_finalize_allreduce_norm, + _spec_rows("m"), +) +MOE_ENGINE = RowRange( + "m3_moe_engine", _rows(1, MOE_ENGINE_MAX_TOKENS), m3_moe_engine, _spec_rows("tokens") +) +MOE16_ENGINE = RowRange( + "m3_moe16_engine", + _rows(MOE16_MIN_TOKENS, MOE16_MAX_TOKENS), + m3_moe16_engine, + _spec_rows("tokens"), +) +# The router of a pure generation step of 9..16 rows when the MoE runs unfused (above the MoE engine's 8; at 9..13 an +# MoE16 engine miss, 14..16 always): the TF32 tensor-core router (minimax_m3_router_tc, 16-CTA clusters, one reduction +# order), the arithmetic the built-in gate takes on SM107 at these rows. Devices that cannot co-schedule 16-CTA +# clusters keep the fp32 router. Context steps never take it. +ROUTER_TC_GEN_MIN_ROWS = 9 # capacity: fixed: past the MoE engine's 8 rows +ROUTER_TC_GEN_MAX_ROWS = 16 # capacity: fixed: the built-in gate's router_tc range +ROUTER_TC_GEN = RowRange( + "m3_router_tc", + _rows(ROUTER_TC_GEN_MIN_ROWS, ROUTER_TC_GEN_MAX_ROWS), + m3_router_tc, + _spec_rows("m"), # every cell checks the logits, with or without the NVFP4 activation +) +# The GEMM paths are keyed (projection (N, K), rows): each projection runs its own split-K / ring / MMA +# configuration, so a row count is certified per projection. +MX_GEMV = RowRange( + "m3_mx_gemv", + frozenset((nk, m) for nk, rows in MX_GEMV_ROWS.items() for m in rows), + m3_mx_gemv, + lambda spec: ((spec["nk"], spec["m"]),), +) +MX_GEMV_BF16 = RowRange( + "m3_mx_gemv_bf16", + frozenset((nk, m) for nk, hi in MX_GEMV_BF16_MAX_TOKENS.items() for m in _rows(1, hi)), + m3_mx_gemv_bf16, + lambda spec: ((spec["nk"], spec["m"]),), +) +# (draft length, requests) of a greedy verify step. +SHARD_ACCEPT = RowRange( + "m3_shard_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SHARD_ACCEPT_MAX_REQUESTS + 1)), + m3_shard_accept, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) + +# Requests of a decode step whose dense-layer attention runs at a capped KV split. +DENSE_DECODE_SPLIT = RowRange( + "m3_dense_decode", + frozenset(DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), + m3_dense_decode, + lambda spec: (len(spec["kv_lens"]),) if spec["capped"] else (), +) + +# (query rows per request, requests) of a decode step whose dense-layer attention runs on PrimTS split-KV decode. +FULLCTX_DECODE = RowRange( + "m3_fullctx_decode", + frozenset(FULLCTX_DECODE_RECIPES) if FULLCTX_DECODE_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + +# (query rows per request, requests) of a generation-only drafter pass that runs on the full-context decode. +DRAFTER_FULLCTX = RowRange( + "m3_drafter_fullctx", + frozenset(FULLCTX_DECODE_RECIPES) if DRAFTER_FULLCTX_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + +# Verify rows of a fused-verify step whose target logits come from the coded head. +TARGET_HEAD = RowRange( + "m3_target_head", + _rows(1, TARGET_HEAD_MAX_ROWS), + m3_target_head, + _spec_rows("m"), +) + +# (draft length, requests) of a fused draft loop, and the rows of its passes' draft input: the first pass's +# requests x (draft length + 1), one per request on the later passes. +DRAFT_STEP = RowRange( + "m3_draft_step", + frozenset((DRAFT_STEP_DRAFT_LEN, n) for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1)), + m3_draft_step, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) +# ((N, K), rows) of a drafter BF16 linear on its low-M kernel. +DRAFT_LOWM = RowRange( + "m3_draft_lowm", + frozenset((nk, m) for nk, rows in DRAFT_LOWM_ROWS.items() for m in rows), + m3_draft_lowm, + lambda spec: ((spec["nk"], spec["m"]),), +) +DRAFT_INPUT = RowRange( + "m3_draft_input", + frozenset( + rows + for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1) + for rows in (n, n * (DRAFT_STEP_DRAFT_LEN + 1)) + ), + m3_draft_input, + _spec_rows("m"), +) + +# The step-loop hooks (RuntimeHooks), one switch each. They move a decode step's input preparation +# off the GPU timeline between step graphs. +STEP_PACK_INPUTS = True +# The pack's host copies go out in one parallel scatter (a program per copy) when their destinations do not overlap: the +# serial one-CTA scatter pays one dependent global round trip per copy (SM107: 12.5 / 15.6 / 31.7 us at 16 / 24 / 40 +# copies vs 6.5 / 6.5 / 6.9 us parallel). +STEP_PACK_PARALLEL_SCATTER = True +# The pack is issued by the decode graph's own prologue (a copy of a pinned window + the parallel scatter), so no +# operation is left between the step graphs (needs STEP_PACK_PARALLEL_SCATTER). +STEP_PACK_IN_GRAPH = True +# The window reaches the prologue's device buffer on a side copy stream while the previous graph runs; the prologue +# waits on that copy and scatters device to device, so the replay no longer starts with a host copy. Needs +# STEP_PACK_IN_GRAPH. +STEP_PACK_EARLY_COPY = True +# The step epoch (step_inputs.step_epoch): one int64 device counter the decode graph's prologue increments once per +# replay (program 0 of its first scatter launch, no extra kernel) and never resets. The per-layer megakernel pieces +# tag their per-step device flags with it (contract v1: tag = epoch << 8 | layer). Needs STEP_PACK_IN_GRAPH. +STEP_PACK_EPOCH = True +# The prologue is one launch of this target's op (ops/spec/m3_step_prologue, minimaxM3Sm107Tp4/stepPrologue.cu): the +# host copies, a CTA barrier, the device copies, the header clear and the epoch, where the default prologue runs two +# Triton scatter launches and an ATen fill (3 graph kernels -> 1, no Triton at the head of the decode graph). Same bytes +# written. Needs STEP_PACK_EARLY_COPY. Kill: False. +STEP_PACK_CPP_PROLOGUE = True +# PRE-WAIT WRITER CONTRACT (rubin-20261015/PREWAIT_AUDIT.md item b): the kernels that read step inputs (seq_lens, +# block tables, positions, cache slots, the sampled-accept pre-wait parameters) before their grid-dependency wait rely +# on every writer of them being a host-to-device copy before the replay or a PLAIN launch earlier in the step: +# msa_kv_update, the step-input prologue (Triton scatter or the C++ op), the spec fold tail (Triton or C++), the pack +# scatter and any step-head op (the spec step correction, STEP_SPEC_OFFSETS_ONE_LAUNCH). Never give one of them a PDL +# attribute with an early trigger. +STEP_SKIP_UNCHANGED_H2D = True +# A changed packed 2-D metadata upload stages only each row's changed column span. On a step where a request crosses +# into a new page the MSA block table changed by one entry, but the whole padded table (32 KB a request) was staged, +# overflowing the 16 KB in-graph window into an eager upload and scatter between the step graphs (57-69 us at the +# ~1 in 19 block-crossing steps; TTFT inter-graph scan). Kill: False. +STEP_H2D_CHANGED_SPAN = True +# A speculative overlap step re-derives the MSA KV-length metadata once, after the forward's KV-length correction, +# instead of also before it (RuntimeHooks.kv_lens_update_once): the first pass (a fill, the out_cache_loc tail fill and +# an msa_kv_update at the head of the decode graph, 2.13 us at TP4) is rewritten in full by the second. Kill: False. +STEP_KV_LENS_UPDATE_ONCE = True +# The overlap scheduler's speculative step correction (position_ids += offsets, kv_lens += offsets) runs as this +# target's one plain launch (RuntimeHooks.spec_offsets_op, ops/spec/m3_spec_offsets) instead of two ATen adds at the +# head of the decode graph: the same int32 adds. TP8 graph bench (the work is per token, not per rank; 64 steps +# per graph, a predecessor kernel each): -1.27 us/step at 4 tokens / 1 request, -0.72..-0.78 at 8-32 tokens. +# Kill: False. +STEP_SPEC_OFFSETS_ONE_LAUNCH = True +STEP_LOOP_ON_EXECUTION_STREAM = True +STEP_MERGE_TUNED_TACTICS = True +STEP_SKIP_GREEDY_RNG_UPLOAD = True +# Sampled speculative batches stage their four per-step Philox seed/offset uploads into the step-input pack +# (STEP_PACK_INPUTS) instead of issuing four copies between the step graphs. +STEP_PACK_SPEC_RNG_UPLOADS = True +# The spec step's small uploads: the accepted-draft counts ride the step-input pack, and the sampler's slot ids stay in +# a device buffer re-uploaded only when they change (a pageable upload every step otherwise). +STEP_SPEC_SMALL_H2D = True +STEP_SPEC_FOLD = True +# The forward of a decode-only step stores its outputs into the spec sampler's store and stages the next step's inputs +# at its tail (one launch inside the step graph), so neither the sampler's store nor the next step's gathers run +# between the step graphs (M3SpecStepFold.in_forward). +STEP_SPEC_FOLD_IN_FORWARD = True +# The forward's tail (STEP_SPEC_FOLD_IN_FORWARD) runs as this target's C++ op (kernels/m3_fold_tail cpp=, +# minimaxM3Sm107Tp4/foldTail.cu) instead of the Triton kernel: the same integer copies, one launch either way (no +# Triton on the generation decode path). Kill: False. +STEP_FOLD_TAIL_CPP = True +# With STEP_SPEC_FOLD_IN_FORWARD the spec sampler's host copies of its store (new tokens, lengths, draft tokens) still +# ran on the main stream between the step graphs (the forward, not the sampler, wrote the store, and it rewrites it at +# the next replay's tail). They run on the D2H side stream instead, and the forward's tail waits on an event the +# sampler records after them before it rewrites the store (M3SpecStepFold.store_copies_done). +STEP_SPEC_STORE_COPY_SIDE = True +# The spec sampler copies its per-slot store to the host on its D2H side stream, so the copies overlap the next +# step's input staging instead of sitting between the step graphs. +STEP_SAMPLER_D2H_SIDE_STREAM = True +# The iteration log times the device step with one event per iteration boundary on a side stream, so no timing +# record sits between two step graphs on the main stream (the reported step times are the same). +STEP_ONE_TIMING_EVENT = True +MSA_SHORT_SEQUENCE_IDENTITY = True +# The sparse decode (the LC, the Triton merge) also writes the MXFP8 o_proj hand-off (bit +# for bit torch.ops.trtllm.mxfp8_quantize of its BF16 rows) and publishes it to m3_prequant.SLOT. Kill: False publishes +# nothing, so the o_proj quantizes its BF16 input itself as the built-in path does. +MERGE_MXFP8_HANDOFF = True +MSA_KV_UPDATE = True +MSA_SKIP_UNCHANGED_BLOCK_OFFSETS = True +MSA_REUSE_POOL_VIEWS = True +# Step-input preparation fills the MSA block-id table by slice-copying each request's page-index view (MSA +# RuntimeHooks.block_ids_from_views) instead of building a Python list of the whole padded row; the table is the same. +MSA_BLOCK_IDS_FROM_VIEWS = True + +FAST_PATHS: Tuple[RowRange, ...] = ( + EMBED_FIRSTNORM, + DENSE_DECODE_SPLIT, + FULLCTX_DECODE, + DRAFTER_FULLCTX, + SPARSE_DECODE_PRODUCER, + PRODUCER, + INDEX_SCORE, + INDEX_SCORE_T2SC_PATH, + LONG_SELECT, + LONG_SELECT_DSL, + MC_ALLREDUCE, + DRAFT_MC_AR, + ATTN_PUSH, + DENSE_PUSH, + MOE_PUSH, + MC_FINALIZE, + MOE_ENGINE, + MOE16_ENGINE, + ROUTER_TC_GEN, + MX_GEMV, + MX_GEMV_BF16, + SHARD_ACCEPT, + TARGET_HEAD, + DRAFT_STEP, + DRAFT_LOWM, + DRAFT_INPUT, +) + +# -------------------------------------------------------------------------------------------------------------------- +# Selection: pure functions of the step shape +# -------------------------------------------------------------------------------------------------------------------- + + +def plan_for(num_tokens: int) -> DecodePlan: + for plan in DECODE_PLANS: + if num_tokens <= plan.max_tokens: + return plan + raise AssertionError(f"no decode plan covers {num_tokens} tokens") + + +def is_decode_step(num_tokens: int) -> bool: + """Whether a pure generation step of ``num_tokens`` rows runs the decode kernels (else generic layers).""" + return num_tokens <= DECODE_MAX_TOKENS + + +def moe_path(num_tokens: int) -> Optional[str]: + """The MoE kernel of a decode step: an engine, the unfinalized MoE + finalize all-reduce, or None (generic).""" + if MOE_ENGINE.covers(num_tokens): + return MOE_ENGINE.path + if MOE16_ENGINE.covers(num_tokens): + return MOE16_ENGINE.path + if MC_FINALIZE.covers(num_tokens): + return MC_FINALIZE.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# The dense layers' swigluoai activation (where the fold above would apply) as this target's C++ op +# (ops/gemm/m3_dense_swiglu, minimaxM3Sm107Tp8/swigluOai.cu: bit for bit the Triton silu_and_mul), PDL-launched with +# an early release so the down GEMV stages its weights during it; the down GEMV reads the activation only after its own +# grid wait. TP2 chain bench: -0.69 us per dense layer vs the Triton kernel. Kill constant: False +# runs the Triton kernel. +DENSE_SWIGLU_CPP = True +DENSE_SWIGLU_CPP_PATH = "m3_dense_swiglu" +# FlashInfer's CuTe DSL MXFP8 GEMM on the producer's MXFP8 activation (rows the GEMV does not serve). +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +# The MXFP8 quantize (tuned backend) and FlashInfer's CuTe DSL GEMM: every row the decode kernels do not serve. +GENERIC_GEMM_PATH = "mxfp8_gemm" + + +def linear_path(n: int, k: int, num_tokens: int, prequant: bool, low_m_gemm: bool) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. ``low_m_gemm`` is accepted for the call sites' + symmetry: no low-M GEMM configuration is certified at TP4, so it selects nothing. + """ + if prequant: + if num_tokens in MX_GEMV_ROWS.get((n, k), ()): + return MX_GEMV_PATH + return PREQUANT_GEMM_PATH + if num_tokens <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0): + return MX_GEMV_BF16_PATH + return GENERIC_GEMM_PATH + + +# -------------------------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A covered decode shape could not take its fast path under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() + + +def set_require(require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: + """The (PDL launch, early trigger, trigger phase) keys of the 384-thread select a decode step can launch, compiled + before warmup; none when ``SELECT_LONG_DSL`` is off.""" + if not LONG_SELECT_DSL.keys: + return () + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + pdl = get_env_enable_pdl() + keys = [(pdl, False, SELECT_LONG_TRIGGER_PHASE), (pdl, False, 0)] + if SELECT_EARLY_MAX_ROWS: + keys.append((pdl, True, 0)) + return tuple(dict.fromkeys(keys)) + + +def missed(path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if _require[0]: + raise FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " + f"{reason}" + ) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def taken(path: str, rows: Optional[int], what: str = "") -> None: + """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line + every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" + at = "" if rows is None else f" at {rows} rows" + detail = f" ({what})" if what else "" + _log_once( + logger.info, + path, + f"taken: {rows}" + (f": {what}" if what else ""), + f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", + ) + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +# Kernel modules report from inside their calls through kernels/m3_report: the same policy. +m3_report.install( + on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed +) + + +# -------------------------------------------------------------------------------------------------------------------- +# Target constants +# -------------------------------------------------------------------------------------------------------------------- + +# The shared expert folds into the routed engine from one token (dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# SMs left to the concurrent kernels while a persistent routed engine runs. +MOE_ENGINE_RESERVED_SMS = 24 # capacity: fixed: an SM split, not a size +# Id of the target's final hidden state in Eagle3's ``layers_to_capture``. +FINAL_HIDDEN_LAYER_ID = -1 # capacity: fixed: an Eagle3 layer id, not a size +# The trtllm-gen MoE runner's routing kind for MiniMax-M3 (sigmoid + correction bias, renormalized, scaled) and its +# gated activation (SwiGLU with the per-expert alpha / beta / clamp operands). +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) + +# -------------------------------------------------------------------------------------------------------------------- +# Streaming: also emit a response on decoding iterations 2..FIRST_CONTENT_EMIT_ITERS (RuntimeHooks), so a request's +# first content leaves one step after its first token. 1: the stock schedule. +FIRST_CONTENT_EMIT_ITERS = 2 + +# The runtime-value dump (TRTLLM_MODELING_V2_VALUE_DUMP=1): the metadata of every fast-path and entry decision a step +# takes, recorded on eager steps and captures (replay runs no Python), printed once at exit. +# -------------------------------------------------------------------------------------------------------------------- + +_VALUE_DUMP = os.environ.get("TRTLLM_MODELING_V2_VALUE_DUMP") == "1" +_VALUES: set = set() + + +def _record(*meta: Hashable) -> None: + if _VALUE_DUMP: + _VALUES.add(meta) + + +def _print_values() -> None: + if _VALUES: + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "0")) + for meta in sorted(_VALUES, key=repr): + print(f"[modeling_v2 value dump] rank {rank}: {meta}", flush=True) + + +if _VALUE_DUMP: + atexit.register(_print_values) + +# -------------------------------------------------------------------------------------------------------------------- +# Configuration: the text tower and the runtime the target assumes + +# -------------------------------------------------------------------------------------------------------------------- + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for (it is 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'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs +# -------------------------------------------------------------------------------------------------------------------- + + +def _mxfp8_linear( + core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + _record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +def _gemv_variants(n: int, k: int, push: Optional[str], world: int): + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch routes to an [N, K] + projection's GEMV: an MXFP8 activation on the rows ``MX_GEMV_ROWS`` gives it, a bf16 one up to + ``MX_GEMV_BF16_MAX_TOKENS``; each pushing when the projection's boundary push covers the rows.""" + push_rows = {"attn": ATTN_PUSH.keys, "dense": DENSE_PUSH.keys, "moe": MOE_PUSH.keys}.get( + push, frozenset() + ) + variants = [] + for rows in range(1, DECODE_MAX_TOKENS + 1): + for x_bf16, served in ( + (False, rows in MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + bucket = next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + variants.append((rows, x_bf16, (world, bucket, AR_PUSH_MODE))) + return variants + + +def _precompile_cfg( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + cfg: tuple, + use_pdl: bool, + ar: Optional[Tuple[int, int, int]] = None, +) -> Optional[tuple]: + """``mx_gemv.precompile`` for a call with this target's (split, ring) override ``cfg`` (``mx_gemv_cfg``): the + shared precompile resolves the shared table's config only. Compiles without launching, outside capture, also while + compiles are frozen; returns the kernel's key, or None when such a call is declined.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = mx_gemv._resolve(m, n, k, False, True, cfg, use_pdl, push) + if resolved is None: + return None + key, runtime_ar = resolved + if mx_gemv.is_compiled(key): + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("_precompile_cfg allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros(m, k, dtype=torch.float8_e4m3fn, device=dev) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = mx_gemv._views(x, xsf, weight, weight_scale, y, False) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + mx_gemv._compile(key, args, m, runtime_ar, stream, mx_gemv._extra_args(dev)) + return key + + +_MX_TAC_ENGAGED: set = set() + + +def _note_mx_tac(kind: str, n: int, k: int, rows: int) -> None: + """MX-TAC's engagement line, once per (activation kind, projection, rows).""" + key = (kind, n, k, rows) + if key not in _MX_TAC_ENGAGED: + _MX_TAC_ENGAGED.add(key) + logger.info( + f"[M3 sm_107 tp4] MX-TAC engaged: pinned MXFP8 decision (quantize tactic " + f"{MXFP8_DECODE_QUANT_TACTIC}, GEMM tactic {MXFP8_DECODE_GEMM_TACTIC}) " + f"for the [{n}, {k}] projection at {rows} rows, {kind} activation" + ) + + +def _uncompiled_note() -> str: + key = mx_gemv.last_uncompiled + return "" if key is None else f" (its kernel {key} was not compiled before warmup)" + + +def _linear( + core, x: torch.Tensor, wk: str, sk: str, low_m_gemm: bool, push: Optional[tuple] = None +) -> Tuple[torch.Tensor, bool]: + """One MXFP8 projection ``x @ w[wk]^T``; returns (output, whether its GEMV pushed into ``push``). + + Outside decode CUDA graphs: the engine-tuned MXFP8 quantize + GEMM. Inside them the projection's (N, K) and the + step's rows pick one path (``linear_path``): + + 1. an activation its producer already quantized to MXFP8 (``m3_prequant.SLOT``, taken here): the tcgen05 decode + GEMV on the rows it serves, else the decode-graph MXFP8 GEMM on that activation (AutoTuner-selected backend + per token bucket); + 2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; + 3. otherwise the engine-tuned MXFP8 quantize + GEMM. + + A kernel declining rows its path covers is a miss; the chain then continues down the list. ``push`` (an + ``m3_mx_gemv`` ``ar`` tuple) is the boundary all-reduce exchange the GEMV writes its rows into. ``low_m_gemm`` + is accepted for the call sites' symmetry (``linear_path``): no TP4 low-M GEMM configuration is built. + """ + w, s = core.w[wk], core.w[sk] + x = x.contiguous() + mode = _mx_mode(core.mxfp8) + if not mode.decode_graph: + return _mxfp8_linear(core, x, w, s, mode), False + use_pdl = get_env_enable_pdl() + rows = x.shape[0] + n, k = w.shape + prequant = m3_prequant.SLOT.take(x) + if prequant is None and x.dtype != torch.bfloat16: + path = GENERIC_GEMM_PATH + else: + path = linear_path(n, k, rows, prequant is not None, low_m_gemm) + _record("linear", (n, k), rows, path, push is not None) + output = None + # MX-TAC: the decode-graph MXFP8 quantize and GEMM at the target's fixed tactics, no AutoTuner. + pinned = MX_TAC_TP4 and core.mxfp8.quant_method.tune_decode_graph_backends + if prequant is not None: + act, act_scale = prequant + if path == MX_GEMV_PATH: + cfg = mx_gemv_cfg(n, k, rows) + output = m3_mx_gemv(act, act_scale, w, s, True, cfg, use_pdl, push) + if output is not None: + split = f", split {cfg[0]}" if cfg is not None else "" + taken(path, rows, f"[{n}, {k}] mxfp8 in{split}") + return output, push is not None and mx_gemv.last_pushed + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if path == PREQUANT_GEMM_PATH: + uncovered(MX_GEMV_PATH, f"the [{n}, {k}] projection at rows the GEMV does not serve") + if pinned: + output = m3_mxfp8_pinned_gemm( + act, act_scale, w, s, torch.bfloat16, MXFP8_DECODE_GEMM_TACTIC + ) + _note_mx_tac("mxfp8", n, k, rows) + return output, False + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16(x, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] bf16 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, rows, f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note() + ) + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + if pinned and x.dtype == torch.bfloat16: + act, act_scale = m3_mxfp8_pinned_quantize(x, MXFP8_DECODE_QUANT_TACTIC) + output = m3_mxfp8_pinned_gemm( + act, act_scale, w, s, torch.bfloat16, MXFP8_DECODE_GEMM_TACTIC + ) + _note_mx_tac("bf16", n, k, rows) + return output, False + # Stock's decode-graph dispatch for this shape (the tuned quantize + GEMM). + return _mxfp8_linear(core, x, w, s, mode), False + + +def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: + """A dense MLP's SwiGLU-OAI activation: this target's C++ op (``DENSE_SWIGLU_CPP``, bit for bit the Triton kernel, + PDL-launched with an early release) up to its certified rows, in any graph or eager step; else the Triton + kernel.""" + rows = int(gu.shape[0]) + path = DENSE_SWIGLU_CPP_PATH + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + if DENSE_SWIGLU_CPP and not is_torch_compiling(): + if rows > max_rows: + uncovered(path, f"the dense swigluoai op: {rows} rows, past its certified {max_rows}") + else: + act = m3_dense_swiglu( + gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + ) + if act is not None: + taken(path, rows) + return act + uncovered(path, "the gate_up layout") + return m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer boundaries: all-reduce + residual + RMSNorm +# -------------------------------------------------------------------------------------------------------------------- + + +def _can_mc(x, folded: bool = False) -> bool: + """Whether the generic boundary takes the multicast all-reduce + residual + RMSNorm (``GENERIC_MC_AR_MAX_TOKENS``, + with the shared + routed add folded in ``GENERIC_MC_AR_FOLDED_MAX_TOKENS``).""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (GENERIC_MC_AR_FOLDED_MAX_TOKENS if folded else GENERIC_MC_AR_MAX_TOKENS) + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + +def _mc_allreduce_norm(core, x, x2, residual, norm_w, *, quant=False, prepushed=0, capture=None): + """The multicast all-reduce + optional second addend + residual + RMSNorm (``m3_mc_allreduce_norm``), or None + without a workspace or for operands the kernel does not take (the same on every rank). With ``quant`` the normed + rows are also published, MXFP8-quantized, to ``m3_prequant.SLOT`` for the next MXFP8 projection; ``capture`` = + (out, fc_norm weight, eps): the boundary also writes the Eagle3 capture norm of the updated residual. Up to + ``MC_AR_CLUSTER_MAX_ROWS`` rows the one-shot 8-CTA cluster form runs; ``MC_AR_CARVEOUT`` sets its shared-memory + carveout.""" + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + if workspace is None: + return None + _record("m3_mc_allreduce_norm", x.shape[0], x2 is not None, quant, prepushed) + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm_w, + core.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + oneshot_cluster=x.shape[0] <= MC_AR_CLUSTER_MAX_ROWS, + carveout=MC_AR_CARVEOUT, + capture=capture, + ) + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm.""" + _record("m3_allreduce", x.shape[0]) + return m3_allreduce( + x, + core.allreduce, + AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ), + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The MoE: router, routed experts (trtllm-gen NVFP4), shared expert +# -------------------------------------------------------------------------------------------------------------------- + + +@dataclass +class _MoeScalars: + """The routed experts' kernel scalars of one layer, derived from the checkpoint's per-tensor scales once the + weights are real (``MiniMaxM3Tp4ModelingV2._derive_moe``).""" + + fc31_input_scale: torch.Tensor + fc2_input_scale: torch.Tensor + fc31_alpha: torch.Tensor + fc2_alpha: torch.Tensor + fc31_scale_c: torch.Tensor + act_alpha: torch.Tensor + act_beta: torch.Tensor + act_clamp: torch.Tensor + + +def _router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The fp32 router logits of MoE layer ``i`` when the MoE runs unfused (the decode steps' engines route + in-kernel): ``F.linear(float(h), w)``, the built-in gate's arithmetic outside its tensor-core router's rows. The + experts quantize their bf16 input themselves.""" + _record("m3_fp32_linear", h.shape[0] if h.dim() == 2 else 0) + return m3_fp32_linear(h, core.w[f"l{i}_router"]) + + +def _decode_router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The router logits of a pure generation step: ``ROUTER_TC_GEN``'s rows take the TF32 tensor-core router when + the device co-schedules 16-CTA clusters (decided at prebuild) and the operands are the built-in gate's (2-D bf16, + unit column stride, 16-byte aligned rows); every other step the fp32 router (``_router_logits``).""" + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + ROUTER_TC_GEN.covers(rows) + and core.router_tc_cluster16 + and h.dtype == torch.bfloat16 + and h.stride(1) == 1 + and (h.stride(0) * 2) % 16 == 0 + ): + taken(ROUTER_TC_GEN.path, rows) + _record("m3_router_tc", rows) + return m3_router_tc(h, core.w[f"l{i}_router"], None, get_env_enable_pdl())[0] + return _router_logits(core, i, h) + + +def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): + """The trtllm-gen NVFP4 MoE runner over this rank's tensor-parallel slice of every expert, routing in-kernel. + + ``x_sf`` None means ``x`` is the bf16 activation, quantized here with the experts' global input scale. Returns + the finalized rows, or the unfinalized ``[gemm2 out, expert scales, expanded->permuted index]``. + """ + moe = core.moe[i] + if x_sf is None: + x, x_sf = fp4_quantize(x, moe.fc31_input_scale, 16, False, False) + rows = x.shape[0] + x_sf = x_sf.view(rows, -1).flatten() + _record("fp4_block_scale_moe_runner", rows, do_finalize) + outs = fp4_block_scale_moe_runner( + router_logits, + core.w[f"l{i}_router_bias"], + x, + x_sf.view(torch.float8_e4m3fn), + core.w[f"l{i}_e_w31"], + core.w[f"l{i}_e_s31"].view(torch.float8_e4m3fn), + None, + moe.act_alpha, + moe.act_beta, + moe.act_clamp, + core.w[f"l{i}_e_w2"], + core.w[f"l{i}_e_s2"].view(torch.float8_e4m3fn), + None, + moe.fc31_scale_c, + moe.fc31_alpha, + moe.fc2_alpha, + core.num_experts, + core.top_k, + None, + None, + core.moe_inter, + 0, + core.num_experts, + core.routed_scaling_factor, + ROUTING_MINIMAX2, + do_finalize, + act_type=ACT_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=core.max_num_tokens, + use_dp=False, + ) + return outs[0] if do_finalize else outs + + +def _split_chunk(rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +# -------------------------------------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------------------------------------- + + +def _nvfp4_layer(manager, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _producer_covered(core) -> bool: + """Whether ``PRODUCER`` covers this target's sparse layers: an NVFP4 KV cache at the TP4 heads (16 query, one KV + and one index head), with the range on. A covered layer that cannot take ``m3_producer`` is a miss.""" + return ( + core.main_kv_is_nvfp4 + and (core.heads, core.kv_heads, core.index_heads) == (16, 1, 1) + and PRODUCER.covers("nvfp4_tp4") + ) + + +def _producer_layout(kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 pools are the layout ``m3_producer`` is built for: [pages, 2, 1, 128, 64] with + its scale pool [pages, 2, 1, 128, 8].""" + return ( + kv_data is not None + and kv_scale is not None + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride()[2:] == (128 * 64, 64, 1) + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride()[2:] == (128 * 8, 8, 1) + and kv_scale.shape[0] == kv_data.shape[0] + ) + + +def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or data.shape[0] != scales.shape[0] + or tuple(data.shape[1:]) != (2, num_heads, 128, 64) + or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) + or data.dtype not in (torch.int8, torch.uint8) + or scales.dtype != torch.uint8 + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or cache.stride(1) < num_heads * cache.stride(2) + or cache.stride(0) < 2 * cache.stride(1) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + + +def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: + if covered: + missed(PRODUCER.path, rows, reason) + return None + + +def _sparse_qkv_producer( + core, i: int, packed, position_ids, kv_cache_manager, msa_idx_k_cache, msa_out_cache_loc +): + """The horizontal producer of a sparse layer: Q / index-Q (FP8) out, K / V / index-K written into the paged + caches (NVFP4 or FP8 main cache), for every step type (token-major). An NVFP4 cache at the TP4 heads takes + ``m3_producer`` (``PRODUCER``); a covered layer that cannot is a miss. Returns (q, index_q), or None for caches + neither kernel takes. + + ``kv_cache_manager``, ``msa_idx_k_cache`` (the per-layer index-K cache getter) and ``msa_out_cache_loc`` are the + calling target's ``step_args`` entries.""" + covered = _producer_covered(core) + rows = int(packed.shape[0]) + if packed.dtype != torch.bfloat16 or position_ids is None: + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + manager = kv_cache_manager + if manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = manager.get_buffers(i, kv_layout="HND") + index_k_cache = msa_idx_k_cache(i) + out_cache_loc = msa_out_cache_loc + layer_uses_nvfp4 = _nvfp4_layer(manager, i) + scale_buffers = ( + manager.get_block_scale_buffers(i, kv_layout="HND") if layer_uses_nvfp4 else None + ) + kv_quant_scale = core.inv_kv_scales[i] if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supported_main = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_fused_write_layout(buffers, scale_buffers, core.kv_heads) + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale.is_cuda + and kv_quant_scale.dtype == torch.float32 + and kv_quant_scale.numel() >= 3 + and kv_quant_scale.is_contiguous() + ) + else: + supported_main = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, core.kv_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= core.kv_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + cos_sin = core.rotary_cos_sin + supported_rope = ( + cos_sin.is_cuda + and cos_sin.dtype == torch.float32 + and cos_sin.is_contiguous() + and cos_sin.dim() == 3 + and tuple(cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main + or not supported_index + or not supported_rope + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < rows + ): + return _producer_miss( + covered and layer_uses_nvfp4, + rows, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the layout the " + "producer is built for", + ) + w = core.w + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + layer_uses_nvfp4 + and PRODUCER.covers("nvfp4_tp4") + and (core.heads, core.kv_heads, core.index_heads) == (16, 1, 1) + and _producer_layout(buffers, scale_buffers) + ): + early_trigger = PRODUCER_EARLY_TRIGGER and rows > PRODUCER_LATE_TRIGGER_MAX_ROWS + _record("m3_producer", rows, early_trigger) + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + core.heads, + get_env_enable_pdl(), + early_trigger, + PRODUCER_PREWAIT_STEP_INPUTS, + ) + # Every token count takes the producer: rows past a decode step log once as one line. + decode_rows = is_decode_step(rows) + taken( + PRODUCER.path, + rows if decode_rows else None, + f"nvfp4_tp4, {'entry' if early_trigger else 'exit'} trigger, " + f"prewait_step_inputs={PRODUCER_PREWAIT_STEP_INPUTS}" + if decode_rows + else "nvfp4_tp4, past the decode rows", + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool {tuple(scale_buffers.shape)} is not the " + "producer's layout", + ) + _record("m3_qkv_insert", rows, layer_uses_nvfp4) + q, index_q = m3_qkv_insert( + packed.contiguous(), + buffers.view(torch.uint8) if layer_uses_nvfp4 else buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.heads, + core.kv_heads, + core.index_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + ) + return q.flatten(1), index_q.flatten(1) + + +def _qk_norm_rope( + core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool, kv_dest: Optional[tuple] = None +) -> torch.Tensor: + """The per-head Gemma RMSNorm of Q and K and the partial NeoX RoPE of a layer's packed [Q | K | V], in one kernel: + in place, or an FP8 copy (Q / K normed and roped, V cast) when ``out_fp8``. A dense layer's FP8 copy carries the + BF16 output's bits cast to E4M3 (``m3_dense_qk_norm_rope_fp8``, ``DENSE_QKV_FP8_VIA_BF16``); with ``kv_dest`` = + (the layer's HND pool, the step's cache slots, live tokens, a [tokens, q heads, 128] FP8 buffer) that kernel also + stores the live tokens' K / V rows into the pool and every row's Q into the buffer.""" + qkv = qkv.contiguous() + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + out_fp8 + and DENSE_QKV_FP8_VIA_BF16 + and i not in core.sparse_layers + and qkv.dtype == torch.bfloat16 + and core.head_dim == 128 + ): + _record("m3_dense_qk_norm_rope_fp8", qkv.shape[0], kv_dest is not None) + return m3_dense_qk_norm_rope_fp8( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + True, + *(kv_dest or ()), + ) + assert kv_dest is None, "only the dense FP8 kernel stores K / V" + args = ( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + 1.0, # factor: no YARN + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + _record("qk_norm_rope", qkv.shape[0], out_fp8) + if out_fp8: + return m3_qk_norm_rope_fp8(*args) + return fused_qk_norm_rope(*args) + + +def _fp8_index_qk_norm_rope( + core, i: int, idx_qk: torch.Tensor, position_ids, msa_idx_k_cache, msa_out_cache_loc +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + _record("m3_fp8_indexer_qk_norm_rope", rows) + return m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer helpers both targets' bodies call: the Eagle3 capture hook and the norm after a layer +# -------------------------------------------------------------------------------------------------------------------- + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + +def _next_norm(core, i: int) -> torch.Tensor: + """The norm the boundary after layer ``i`` applies: the next layer's input norm, the final norm after the last.""" + return core.w[f"l{i + 1}_in_norm"] if i + 1 < core.num_layers else core.w["final_norm"] + + +# -------------------------------------------------------------------------------------------------------------------- +# The core +# -------------------------------------------------------------------------------------------------------------------- + +# Engine states (compiled kernel + workspace) per (device, token count[, shared fold]), shared by every layer: a step +# runs the layers one after another on one stream. +_ENGINE_STATES: Dict[tuple, object] = {} +_MOE16_STATES: Dict[tuple, object] = {} + + +class MiniMaxM3Tp4ModelingV2(DecoderModel): + """The MiniMax-M3 text model of this target: its weights, its runtime resources, and the per-step dispatch.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_runtime(model_config) + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + self.mapping = mapping + self.vocab_size = cfg.vocab_size + d = _weights.MODEL_WEIGHTS.dims(self) + self.num_layers = d.num_layers + self.hidden = d.hidden + self.heads = d.heads + self.kv_heads = d.kv_heads + self.head_dim = d.head_dim + self.index_heads = d.index_heads + self.index_dim = d.index_dim + self.q_size = d.heads * d.head_dim + self.kv_size = d.kv_heads * d.head_dim + self.moe_inter = d.moe_inter + self.num_experts = d.num_experts + self.top_k = int(cfg.num_experts_per_tok) + self.routed_scaling_factor = float(getattr(cfg, "routed_scaling_factor", 1.0)) + self.rms_eps = float(cfg.rms_norm_eps) + self.swiglu_alpha = float(getattr(cfg, "swiglu_alpha", 1.702)) + self.swiglu_beta = 1.0 # the (up + 1) offset of SwiGLU-OAI + self.swiglu_limit = float(getattr(cfg, "swiglu_limit", 7.0)) + self.max_num_tokens = model_config.max_num_tokens + self.moe_max_num_tokens = model_config.moe_max_num_tokens + self.sparse_layers = frozenset(_sparse_layer_ids(cfg)) + self.moe_layers = frozenset(_moe_layer_ids(cfg)) + mode = model_config.quant_config.quant_mode + self.main_kv_is_fp8 = bool(mode is not None and mode.has_fp8_kv_cache()) + self.main_kv_is_nvfp4 = bool(mode is not None and mode.has_fp4_kv_cache()) + # The sparse layers' QKV and index projections: one five-way packed projection, or the QKV projection and a + # separate BF16 index projection (fuse_qkv_index_projection=False) run on two streams. + self.fused_projection = bool( + getattr(model_config.sparse_attention_config, "fuse_qkv_index_projection", False) + ) + + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # The engine reads these two off the module tree: the MXFP8 GEMM state it tunes, and the opt-in to the + # automatic FlashInfer MXFP8 path on decode graphs. + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + + # The built-in fused all-reduce + residual + RMSNorm of the layer boundaries the multicast kernel does not take. + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + + # Streams: the shared / routed expert overlap and, when the MoE chunks its input, the chunking overlap. + self.aux_stream_moe = torch.cuda.Stream() + chunked = self.moe_max_num_tokens < model_config.max_num_tokens * mapping.dp_size + self.aux_stream_moe_chunk = torch.cuda.Stream() if chunked else None + self.moe_events = {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers} + self.moe_chunk_events = { + i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers + } + # The unfused sparse layers' index branch overlaps the QKV branch on its own stream. + self.aux_stream_attn = None if self.fused_projection else torch.cuda.Stream() + self.attn_events = ( + {} + if self.fused_projection + else {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.sparse_layers} + ) + + # Partial NeoX RoPE (rotary_dim of head_dim): the cos/sin table every attention kernel reads. + rope = RopeParams.from_config(cfg) + self.rope_theta = rope.theta + self.rotary_dim = int(rope.dim) + self.rotary_cos_sin = rope.create_rope_const_params(interleave=False)[1].reshape( + rope.max_positions, 2, -1 + ) + + # One MSA attention backend per layer (dense and sparse; the step metadata is the backend's). + quant_config = model_config.get_quant_config() + self.attn: List[MiniMaxM3MsaSparseAttention] = [] + for i in range(self.num_layers): + backend = create_attention( + model_config.attn_backend, + i, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=None, + quant_config=quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + q_scaling=1.0, + sparse_params=model_config.sparse_attention_config.to_sparse_params( + pretrained_config=cfg, layer_idx=i + ), + ) + backend.update_quant_config(quant_config) + if not isinstance(backend, MiniMaxM3MsaSparseAttention): + raise ValueError( + f"layer {i}: expected the MiniMax-M3 MSA backend, got {type(backend).__name__}" + ) + if i not in self.sparse_layers: + # The KV split trtllm-gen dense decode runs at, per request count. + backend.dense_decode_max_seq_len_by_batch = DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + self.attn.append(backend) + + # Built once the weights are real (``post_load_weights``): the KV (Q, K, V) scales, the MoE scalars, the + # decode engines and workspaces, the two targets. + self.kv_scales: Dict[int, torch.Tensor] = {} + self.inv_kv_scales: Dict[int, torch.Tensor] = {} + self.moe: Dict[int, _MoeScalars] = {} + self.mx_alpha: Optional[torch.Tensor] = None + self.moe_engines: Dict[int, Dict[Tuple[int, bool, bool], Optional[object]]] = {} + self.moe16_engines: Dict[int, Dict[int, Optional[object]]] = {} + self.embed_owners: Dict[tuple, DecodeEmbedding] = {} + # Whether the 9..16-row generation router (ROUTER_TC_GEN) runs; decided at prebuild. + self.router_tc_cluster16 = False + # Each sparse layer's selected block ids per query row at its last long-context decode, read by the previous + # sparse layer's decode to prefetch this layer's KV (``SPARSE_DECODE_LC_PREFETCH_NEXT_KV``). + self.n1_ids: Dict[int, torch.Tensor] = ( + { + i: torch.full( + (N1_MAX_ROWS, 16), + -1, + dtype=torch.int32, + device=torch.device("cuda", torch.cuda.current_device()), + ) + for i in sorted(self.sparse_layers) + } + if SPARSE_DECODE_LC_PREFETCH_NEXT_KV + else {} + ) + self._built_for: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def sm_scale(self, i: int) -> float: + return (self.head_dim**-0.5) / float(self.attn[i].q_scaling) + + # ---------------------------------------------------------------------------------------------- after the load + + def post_load_weights(self) -> None: + """Derive the kernels' state from the loaded weights, then build both targets and every decode kernel. + + Idempotent per set of weight storages: the loader's post-load walk reaches this module directly and through + the shell, and a second call on the same storages does nothing. + """ + identity = ( + self.w["embed"].data_ptr(), + self.w["l0_in_norm"].data_ptr(), + self.w["final_norm"].data_ptr(), + ) + if self._built_for == identity: + return + device = self.w["embed"].device + derived = nn.Module() + derived.register_buffer( + "mx_alpha", torch.ones(1, dtype=torch.float32, device=device), persistent=False + ) + for i in range(self.num_layers): + # No KV calibration in the checkpoint: unit (Q, K, V) scales and their reciprocals. + derived.register_buffer( + f"l{i}_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + derived.register_buffer( + f"l{i}_inv_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + for i in sorted(self.moe_layers): + for name, value in dataclasses.asdict(self._derive_moe(i)).items(): + derived.register_buffer(f"l{i}_{name}", value, persistent=False) + self.derived = derived + self.mx_alpha = derived.mx_alpha + self.kv_scales = {i: getattr(derived, f"l{i}_kv_scales") for i in range(self.num_layers)} + self.inv_kv_scales = { + i: getattr(derived, f"l{i}_inv_kv_scales") for i in range(self.num_layers) + } + self.moe = { + i: _MoeScalars( + **{ + f.name: getattr(derived, f"l{i}_{f.name}") + for f in dataclasses.fields(_MoeScalars) + } + ) + for i in sorted(self.moe_layers) + } + self._build_embed_owners() + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self._built_for = identity + + def cache_derived_state(self) -> None: + """The staged loaders' entry to the same build (``post_load_weights``).""" + self.post_load_weights() + + def _derive_moe(self, i: int) -> _MoeScalars: + """The routed experts' kernel scalars from the checkpoint's per-tensor scales (trtllm-gen NVFP4 layout). + + Global input scales: the reciprocal of the largest per-expert input scale (gate and up must agree). Per expert: + ``alpha = 1 / (input_scale * (1 / weight_scale_2))`` for the gate (fc31_alpha), the up half (its GEMM2-input + rescale ``fc31_scale_c = fc2_input_scale * up_alpha``) and the down projection; SwiGLU-OAI's activation + constants per expert, with beta and the clamp in the pre-dequant GEMM output domain (divided by fc31_alpha). + """ + w = self.w + isc1, isc3, isc2 = w[f"l{i}_e_isc1"].data, w[f"l{i}_e_isc3"].data, w[f"l{i}_e_isc2"].data + if not torch.allclose(isc1, isc3): + raise ValueError(f"layer {i}: the routed experts' w1 and w3 input scales differ") + fc31_input_scale = isc1.max().reciprocal() + fc2_input_scale = isc2.max().reciprocal() + ws2_1, ws2_3, ws2_2 = ( + w[f"l{i}_e_ws2_1"].data, + w[f"l{i}_e_ws2_3"].data, + w[f"l{i}_e_ws2_2"].data, + ) + fc31_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_1)) + up_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_3)) + fc2_alpha = 1.0 / (fc2_input_scale * (1.0 / ws2_2)) + fc31_scale_c = fc2_input_scale * up_alpha + n = self.num_experts + device = fc31_alpha.device + act_alpha = torch.full((n,), self.swiglu_alpha, dtype=torch.float32, device=device) + act_beta = torch.full((n,), self.swiglu_beta, dtype=torch.float32, device=device) + act_clamp = torch.full((n,), self.swiglu_limit, dtype=torch.float32, device=device) + act_beta.div_(fc31_alpha) + act_clamp.div_(fc31_alpha) + return _MoeScalars( + fc31_input_scale=fc31_input_scale, + fc2_input_scale=fc2_input_scale, + fc31_alpha=fc31_alpha, + fc2_alpha=fc2_alpha, + fc31_scale_c=fc31_scale_c, + act_alpha=act_alpha, + act_beta=act_beta, + act_clamp=act_clamp, + ) + + def _build_embed_owners(self) -> None: + """The fused embedding + first RMSNorm's caller-owned buffers, one per (rows, id dtype) of a decode step.""" + table, weight = self.w["embed"], self.w["l0_in_norm"] + self.embed_owners = { + (rows, dtype): DecodeEmbedding(table, weight, self.rms_eps, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + + def projections(self) -> List[Tuple[str, str, Optional[str]]]: + """Every MXFP8 projection in layer order: (weight key, scale key, the boundary its GEMV pushes into).""" + out = [] + for i in range(self.num_layers): + out.append((f"l{i}_qkv_w", f"l{i}_qkv_s", None)) + out.append((f"l{i}_o_w", f"l{i}_o_s", "attn")) + if i in self.moe_layers: + out.append((f"l{i}_sh_gu_w", f"l{i}_sh_gu_s", None)) + out.append((f"l{i}_sh_dn_w", f"l{i}_sh_dn_s", "moe")) + else: + out.append((f"l{i}_mlp_gu_w", f"l{i}_mlp_gu_s", None)) + out.append((f"l{i}_mlp_dn_w", f"l{i}_mlp_dn_s", "dense")) + return out + + def decode_query_lens(self) -> List[int]: + """The query rows per request a decode step can have: 1, and 1 + each draft length speculation runs.""" + lens = {1} + spec = getattr(self.model_config, "spec_config", None) + if spec is not None: + for drafts in ( + getattr(spec, "max_draft_len", None), + getattr(spec, "max_total_draft_tokens", None), + *(getattr(spec, "draft_len_schedule", None) or {}).values(), + ): + if drafts: + lens.add(1 + int(drafts)) + return sorted(lens) + + def prebuild_kernels(self) -> Dict[str, object]: + """Compile (and load) every kernel a decode forward dispatches to that would otherwise compile on its first + call: each projection's decode GEMVs (every row count, activation and boundary push), the MoE engines of + ``MOE_ENGINE`` / ``MOE16_ENGINE``'s token counts, and the index scorer at the decode query lengths. Runs on the + current (main) stream before the executor's warmup, in the same order on every rank, then synchronizes; after + it, GEMV compiles are frozen and the forward only looks kernels up (a missing one is a miss). It also builds + the T2SC scorer (``INDEX_SCORE_T2SC_PATH``), the 384-thread long-row select (``LONG_SELECT_DSL``) and the dense + layers' full-context decode plans (``FULLCTX_DECODE``). Returns the seconds each part took and the scorer + keys.""" + times: Dict[str, object] = {} + # The 9..16-row generation router: decided once, from the device, never inside a CUDA-graph capture. + # The built-in gate's conditions, mirrored exactly: 128 experts, a hidden size in 1024s, 16-CTA clusters. + self.router_tc_cluster16 = ( + self.num_experts == 128 + and self.hidden % 1024 == 0 + and bool(torch.ops.trtllm.minimax_m3_router_tc_cluster16(self.hidden)) + ) + logger.info( + f"[M3 prebuild] rank {mpi_rank()}: TF32 tensor-core router for {ROUTER_TC_GEN_MIN_ROWS}.." + f"{ROUTER_TC_GEN_MAX_ROWS}-row generation steps: " + + ( + "engaged (16-CTA clusters)" + if self.router_tc_cluster16 + else "not used (no 16-CTA clusters)" + ) + ) + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + if PREPROCESS_TOKMAP_FIX: + # Every path that can run the stock generation preprocess (the drafter's full-context decode or its + # fallback, any trtllm-gen fallback), so it is installed whatever else is enabled. + _tokmap_guard.install( + torch.device("cuda", torch.cuda.current_device()), + DECODE_MAX_TOKENS, + lambda rows: taken(PREPROCESS_TOKMAP_PATH, rows, "one-request token map"), + lambda rows, reason: missed(PREPROCESS_TOKMAP_PATH, rows, reason), + ) + mx_gemv.freeze(False) + use_pdl = get_env_enable_pdl() + world = int(self.mapping.tp_size) + start = time.perf_counter() + for wk, sk, push in self.projections(): + weight, scale = self.w[wk], self.w[sk] + mx_gemv.ensure_compiled(weight, scale, use_pdl) + n, k = weight.shape + for rows, x_bf16, ar in _gemv_variants(n, k, push, world): + # An MXFP8 call this target gives a (split, ring) override (``mx_gemv_cfg``, as ``_linear`` passes + # it) compiles that config. + cfg = None if x_bf16 else mx_gemv_cfg(n, k, rows) + if cfg is not None: + _precompile_cfg(weight, scale, rows, cfg, use_pdl, ar=ar) + else: + mx_gemv.precompile(weight, scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + torch.cuda.synchronize() + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for i in sorted(self.moe_layers): + self._prebuild_engines(i) + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = self._prebuild_index_scorer( + self.decode_query_lens(), self.w["embed"].device + ) + times["scorer_t2sc_keys"] = self._prebuild_index_scorer_t2sc( + self.decode_query_lens(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["scorer_s"] = time.perf_counter() - start + start = time.perf_counter() + times["select_keys"] = [] + if long_select_dsl_keys(): + # The 384-thread long-row select (LONG_SELECT_DSL): one kernel per launch key. + times["select_keys"] = _select_long_dsl_kernel.prebuild( + long_select_dsl_keys(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["select_s"] = time.perf_counter() - start + start = time.perf_counter() + times["fullctx_plans"] = 0 + if FULLCTX_DECODE.keys: + # The dense layers' full-context decode: one compiled PrimTS plan per covered (rows, requests). + times["fullctx_plans"] = _fullctx_decode_kernel.prebuild( + self.w["embed"].device, + {key: FULLCTX_DECODE_RECIPES[key] for key in FULLCTX_DECODE.keys}, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), + get_env_enable_pdl(), + ) + torch.cuda.synchronize() + times["fullctx_s"] = time.perf_counter() - start + mx_gemv.freeze(True) + logger.info( + "[M3 prebuild] decode kernels built before warmup: GEMVs %.1f s, MoE engines %.1f s, index scorer " + "%.1f s (%d variants), long-row select %.1f s (%d variants), full-context decode %.1f s (%d plans)" + % ( + times["gemv_s"], + times["moe_s"], + times["scorer_s"], + len(times["scorer_keys"]), + times["select_s"], + len(times["select_keys"]), + times["fullctx_s"], + times["fullctx_plans"], + ) + ) + return times + + # ------------------------------------------------------------------------------------------ the MoE engines + + def _expert_operands(self, i: int): + """The routed experts' operands the engines read (trtllm-gen NVFP4 SwiGLU-OAI layout), and the input scale.""" + w, m = self.w, self.moe[i] + tensors = [ + w[f"l{i}_e_w31"].data, + w[f"l{i}_e_s31"].data, + w[f"l{i}_e_w2"].data, + w[f"l{i}_e_s2"].data, + m.fc31_alpha, + m.fc31_scale_c, + m.act_alpha, + m.act_beta, + m.act_clamp, + m.fc2_alpha, + ] + return tensors, m.fc31_input_scale + + def _shared_weights(self, i: int): + w = self.w + return ( + w[f"l{i}_sh_gu_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_gu_s"].data, + w[f"l{i}_sh_dn_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_dn_s"].data, + ) + + def engine(self, i: int, num_tokens: int, shared: bool) -> Optional[object]: + """Layer ``i``'s prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in (and the shared + expert with ``shared``). None after a miss: the build failed or was never made (a forward never builds one).""" + return self._lookup( + i, self.moe_engines[i], (num_tokens, True, shared), MOE_ENGINE.path, num_tokens + ) + + def moe16_engine(self, i: int, num_tokens: int) -> Optional[object]: + """Layer ``i``'s prebuilt 9..13-row engine (router, routed and shared experts in one kernel), or None.""" + if not MOE16_MIN_TOKENS <= num_tokens <= MOE16_MAX_TOKENS: + return None + return self._lookup(i, self.moe16_engines[i], num_tokens, MOE16_ENGINE.path, num_tokens) + + def _lookup(self, i: int, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {i}'s engine" + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, what) + else: + missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + missed(path, num_tokens, "its build failed (logged)") + return engine + + def _prebuild_engines(self, i: int) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of layer ``i`` can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in, and ``MOE16_ENGINE``'s. Whether an + engine is kept is agreed across the TP group (one host collective per engine): a build that fails on one rank + drops that engine on every rank, so the ranks never mix an engine (and its boundary push) with the generic + MoE.""" + engines = self.moe_engines.setdefault(i, {}) + engines16 = self.moe16_engines.setdefault(i, {}) + for num_tokens in sorted(MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in engines: + engines[key] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build( + i, MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + f"[M3 dispatch] {MOE_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + for num_tokens in sorted(MOE16_ENGINE.keys): + if num_tokens not in engines16: + engines16[num_tokens] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build(i, MOE16_ENGINE.path, num_tokens, self._build_moe16_engine), + f"[M3 dispatch] {MOE16_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + + def _try_build(self, i: int, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(i, num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {i}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, i: int, num_tokens: int, shared: bool) -> object: + tensors, scale = self._expert_operands(i) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + kwargs["shared_weights"] = self._shared_weights(i) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = M3MoeState( + device, num_tokens=num_tokens, num_ctas=sms - MOE_ENGINE_RESERVED_SMS, shared=shared + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + def _build_moe16_engine(self, i: int, num_tokens: int) -> object: + tensors, scale = self._expert_operands(i) + device = tensors[0].device + state = _MOE16_STATES.get((str(device), num_tokens)) + if state is None: + state = _MOE16_STATES[(str(device), num_tokens)] = M3Moe16State( + device, num_tokens=num_tokens + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + shared_weights=self._shared_weights(i), + ) + # Compile (first use of the state) and warm. + m3_moe16_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + # ------------------------------------------------------------------------------------------ the index scorer + + def _scorer_variants(self, query_lens) -> list: + """The scorer keys a decode step can take: every sparse layer's (index dtype, index heads per rank, index dim) + at each decode query length, with the early-issue schedule and without it, where the kernel serves it.""" + runner = _index_score_kernel.IndexScoreRunner + keys = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=config.num_index_heads, + head_dim=config.sparse_index_dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for early in (True, False): + key = ( + dtype, + int(config.num_index_heads), + int(qlen), + int(config.sparse_index_dim), + int(INDEX_SCORE_SPLIT_K), + bool(early), + ) + if key not in keys: + keys.append(key) + return keys + + def _prebuild_index_scorer(self, query_lens, device) -> list: + """Compile every scorer key and launch each once on real-shaped dummy operands (loading its module), on the + current stream. The kernel's batch, token, page and block counts are symbolic: its compile needs no step + metadata. Returns the keys.""" + runner = _index_score_kernel.IndexScoreRunner + keys = self._scorer_variants(query_lens) + page = runner.SUPPORTED_PAGE_SIZE + for key in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + rows = 1 if early else INDEX_SCORE_EARLY_MAX_REQUESTS + 1 + blocks = 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros( + rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device + ).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + runner.kernel_cache[key]( + idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + ) + return keys + + def _prebuild_index_scorer_t2sc(self, query_lens, device) -> list: + """T2SC (``INDEX_SCORE_T2SC_PATH``): compile this target's copy of the TP2 / TP8 scorer at every covered request + count's schedule (``index_score_t2sc_schedule``) for every sparse layer geometry and decode query length, and + launch each once on real-shaped dummy operands (loading its module). Returns the keys.""" + runner = _index_score_t2sc_kernel.IndexScoreRunner + keys = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=config.num_index_heads, + head_dim=config.sparse_index_dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for requests in sorted(INDEX_SCORE_T2SC_PATH.keys): + split_k, early = index_score_t2sc_schedule(requests) + key = ( + dtype, + int(config.num_index_heads), + int(qlen), + int(config.sparse_index_dim), + int(split_k), + bool(early), + ) + if all(key != seen for seen, _ in keys): + keys.append((key, requests)) + page = runner.SUPPORTED_PAGE_SIZE + for key, rows in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + blocks = 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros( + rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device + ).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + runner.kernel_cache[key]( + idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + ) + return [key for key, _ in keys] + + # --------------------------------------------------------------------------------------------- the dispatch + + def _select_target(self, attn_metadata: AttentionMetadata): + """Which target runs this step: one host-side integer, fixed per CUDA-graph capture. A step with context rows + (a mixed step included) is prefill's; every pure generation step is decode's, whatever its width.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + capture: Callable[..., None] = _plain_capture, + capture_target: Optional[Callable[..., Any]] = None, + **kwargs, + ) -> torch.Tensor: + """``capture_target``: the worker's capture-norm target (``CAPTURE_NORM_IN_BOUNDARY``) for a capture layer + on the decode fast path; None: the capture runs its own norm.""" + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if kwargs.get("lora_params") is not None: + raise ValueError("this MiniMax-M3 target serves no LoRA adapters") + return self._select_target(attn_metadata).forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The targets +# -------------------------------------------------------------------------------------------------------------------- + + +class PrefillTarget: + """A step with context rows: prefill chunks, and the mixed steps that carry generation rows beside them. Every + layer runs in its generic form; the attention core is the MSA backend's for every row.""" + + def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + Counts and per-step host scalars are read here; device tensors are the metadata's persistent buffers, held + by reference; ``msa_idx_k_cache`` is the metadata's per-layer index-K cache getter. ``metadata`` is the + prepared object itself: the MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_tokens=md.num_tokens, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + """``capture_target`` is unused: the generic boundary takes no multicast all-reduce at this TP + (``GENERIC_MC_AR_MAX_TOKENS`` 0), so a capture layer always runs its own capture norm.""" + core = self.core + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) prefill CUDA graphs" + ) + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + _record("prefill", int(step["num_tokens"]), int(step["num_contexts"])) + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + hidden, residual = self._layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + return hidden + + def _swiglu_mlp( + self, h: torch.Tensor, prefix: str, low_m_gemm: bool, dense: bool = False + ) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP with ``dense``, the shared expert): gate_up, SwiGLU-OAI, down, no + reduction.""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + out, _ = _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm) + return out + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: + the MoE scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: the fp32 router logits and the experts quantizing their bf16 input + (chunked).""" + return self._routed_chunks(i, h, _router_logits(self.core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool): + """The MoE block of layer ``i`` without its reduction. Returns ``(out, routed)``: with ``return_parts`` the + shared and routed parts (routed None when the engine folded the shared expert in), else their sum in + ``out``.""" + core = self.core + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, False) + folded = m3_moe_engine(h.contiguous(), engine, None, None) + return folded, None + routed_out, shared_out = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the + # same collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step): + """Attention of layer ``i`` up to its o_proj (no reduction). Returns (o_proj output, pushed): a prefill step + arms no boundary push, so pushed is always False.""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v for an FP8 cache, else in place in bf16), the + # MSA backend (it writes the layer's cache), o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection + and its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and + the FP8 indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the + o_proj. A layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): the MSA + backend (``m3_msa_attention``: the KV write unless a producer wrote the caches, the sparse layers' top-k block + selection, the FMHA over the context rows and any generation rows beside them).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + sparse = i in core.sparse_layers + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + +class DecodeTarget: + """A pure generation step (no context rows), of any width. + + Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids, per layer: attention (o_proj all-reduce deferred, its rows + pushed into the boundary all-reduce up to 16 tokens) -> attention boundary (all-reduce + residual + post-attention + RMSNorm) -> MoE (the 1..8 / 9..13-row engines, the first pushing its sum into the boundary all-reduce; the + unfinalized MoE with the finalize in the all-reduce from 14 rows) or dense MLP -> MoE boundary (all-reduce + + shared/routed add + residual + the next layer's RMSNorm, emitting the next QKV's MXFP8 activation). A layer whose + output Eagle3 captures runs the generic layer, as does the first layer when the fused embedding does not take the + step. A wider step (or one carrying embeddings) runs every layer in its generic form, with this target's attention + (the decode selection and sparse decode where the producer wrote the caches). + """ + + def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # Set by a fused feed-forward boundary that wrote the layer's Eagle3 capture (see _layer's ``capture``). + self._capture_written = False + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + States the two context counts instead of reading them back: this target runs only when ``num_contexts`` is + zero, so reading them would ask the metadata to confirm what the routing already decided. + + Graph classes: counts, the generation span, the identity-selection flag and the worst-case score extent are + host values fixed per capture (prepare() stages them before the forward); the block table, lengths, cache + slots, identity table and valid-block counts are the metadata's persistent buffers, held by reference and + refreshed in place every step; ``msa_idx_k_cache`` / ``msa_proxy_max_score_view`` are the metadata's getters + of the per-layer index-K cache and of the score scratch view. ``metadata`` is the prepared object itself: the + MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=0, + num_ctx_tokens=0, + num_tokens=md.num_tokens, + num_seqs=md.num_seqs, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + msa_decode_span=md.msa_decode_span, + msa_block_table=md.msa_block_table, + msa_seq_lens_cuda=md.msa_seq_lens_cuda, + msa_short_ok=md._msa_short_ok, + msa_short_table=md.msa_short_table, + msa_proxy_max_score_view=md.msa_proxy_max_score_view, + msa_worst_case_max_k_tiles=md.msa_worst_case_max_k_tiles, + msa_n_valid_blocks=md.msa_n_valid_blocks, + max_seq_len=md.max_seq_len, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + # The decode kernels' rows (a per-capture constant): up to DECODE_MAX_TOKENS tokens of 1-D input ids. + kernels = ( + inputs_embeds is None and input_ids.ndim == 1 and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + _record("decode", num_tokens, plan.name) + hidden, residual = self._embed(input_ids, num_tokens) + else: + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + _record("decode", int(step["num_tokens"]), "generic") + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + is_capture = spec_metadata is not None and spec_metadata.is_layer_capture(i) + if ( + kernels + and residual is not None + and is_capture + and CAPTURE_LAYERS_FAST + and CAPTURE_FUSED_BOUNDARY + ): + # The capture layer on the fast path (its boundary pushes); the capture reads the MoE / dense + # boundary's reduced residual sum. With CAPTURE_NORM_IN_BOUNDARY the boundary all-reduce writes the + # capture norm itself. + capture_out = ( + capture_target(spec_metadata, i, num_tokens) + if capture_target is not None and CAPTURE_NORM_IN_BOUNDARY + else None + ) + self._capture_written = False + hidden, residual = self._layer( + i, plan, position_ids, hidden, residual, step, capture=capture_out + ) + taken("capture_layer_fast", num_tokens, f"layer {i}") + if capture_out is not None and self._capture_written: + taken("capture_norm_in_boundary", num_tokens, f"layer {i}") + capture(spec_metadata, i, residual, None, written=True) + else: + capture(spec_metadata, i, residual, None) + elif not kernels or residual is None or is_capture: + hidden, residual = self._generic_layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + else: + hidden, residual = self._layer(i, plan, position_ids, hidden, residual, step) + return hidden + + def _embed(self, input_ids: torch.Tensor, num_tokens: int): + """(hidden, residual) of the fused embedding + first RMSNorm; (embedding, None) when it does not run, so the + first layer runs its own input norm.""" + core = self.core + if not EMBED_FIRSTNORM.covers(num_tokens): + uncovered(EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens") + return m3_embedding(input_ids, core.w["embed"]), None + owner = core.embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + missed( + EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 after the load)" + if owner is None + else "the input ids are not contiguous", + ) + return m3_embedding(input_ids, core.w["embed"]), None + _record("m3_embedding_firstnorm", num_tokens, input_ids.dtype) + taken(EMBED_FIRSTNORM.path, num_tokens, str(input_ids.dtype)) + return m3_embedding_firstnorm( + input_ids, owner.table, owner.weight, owner.eps, owner.raw, owner.output, owner.use_pdl + ) + + def _bucket_flags(self, num_tokens: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(self.core.mapping, self.core.hidden) + workspace = None if workspace is None else m3_mc_allreduce._bucket(workspace, num_tokens) + flags = None if workspace is None else workspace.flags_for(8) + return None if flags is None else (workspace, flags) + + @staticmethod + def _no_bucket(path: str, num_tokens: int) -> None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, num_tokens: int): + """The o_proj's push tuple into the attention boundary's all-reduce, or None (after a miss). + + None comes only from the token count or from workspace state created collectively in the same eager call on + every rank, so all ranks agree. (A lone non-pushing rank would still be safe: its all-reduce pushes the same + bf16 rows into the same set A.)""" + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(ATTN_PUSH.path, num_tokens) + return None + workspace, flags = found + return ( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + AR_PUSH_MODE, + ) + + def _moe_push_targets(self, num_tokens: int): + """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None + (after a miss). + + This decision must match across ranks: prepushed 3 reads set B of every rank, which a non-pushing rank never + fills. None comes only from the token count or from state created in the same eager call on every rank + (workspace, flags, the descriptor cache), and the engine it is armed for is agreed at build time + (``_prebuild_engines``), so it does.""" + hidden = self.core.hidden + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(MOE_PUSH.path, num_tokens) + return None + workspace, flags = found + key = (workspace.mc.data_ptr(), flags.data_ptr()) + desc = self._moe_push_descs.get(key) + if desc is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor") + return None + desc = self._moe_push_descs[key] = ar_push_desc( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + flags.device, + ) + set_b = workspace.mc.data_ptr() + 2 * 4 * workspace.max_tokens * workspace.world_size * ( + hidden // 2 + ) + return desc, ( + set_b, + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + 2, + ) + + def _boundary(self, use_mc: bool, x, residual, norm_w, x2=None, quant=False, capture=None): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce. + ``capture``: the Eagle3 capture the multicast kernel also writes (``_capture_written`` records it).""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant, capture=capture) + if fused is not None: + if capture is not None: + self._capture_written = True + rows = x.shape[0] + taken( + MC_ALLREDUCE.path, + rows, + ("x + x2" if x2 is not None else "x") + + (", MXFP8 copy" if quant else "") + + (", cluster" if rows <= MC_AR_CLUSTER_MAX_ROWS else ""), + ) + return fused + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + missed( + MC_ALLREDUCE.path, + x.shape[0], + m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + ) + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _prepushed(self, i: int, x, residual, norm_w, quant, prepushed, x2=None, capture=None): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed, capture=capture + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + if capture is not None: + self._capture_written = True + rows = x.shape[0] + taken( + MC_ALLREDUCE.path, + rows, + f"prepushed {prepushed}" + (", cluster" if rows <= MC_AR_CLUSTER_MAX_ROWS else ""), + ) + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step, capture=None): + """``capture``: (out, fc_norm weight, eps) for the feed-forward boundary to write the layer's Eagle3 capture + norm (CAPTURE_NORM_IN_BOUNDARY, for a capture layer on this path); a fused boundary that writes it sets + ``_capture_written``.""" + core = self.core + num_tokens = hidden.shape[0] + # Padded rows: the boundary all-reduce polls and re-arms rows t < M = residual.shape[0], CUDA-graph padding + # rows included. Every producer pushes all rows of the same tensor: the o_proj / shared down GEMV t < m = + # x.shape[0] (its row count comes from the host shape, not a live count), the MoE engine every row of the + # num_tokens it was built for (= num_tokens here). The all-reduce checks x.sizes() == residual.sizes(), so + # no padded row is left with a stale flag. + # The GEMV dispatch (and with it every push into a boundary all-reduce) is a decode-graph one. + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + uncovered( + "decode_graph", + "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + ) + push = plan.attn_mc_ar and decode_graph and ATTN_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if push else None + hidden, attn_pushed = self._attention(i, position_ids, hidden, step, armed) + post_norm = core.w[f"l{i}_post_norm"] + if armed is not None and attn_pushed: + taken(ATTN_PUSH.path, num_tokens) + hidden, residual = self._prepushed(i, hidden, residual, post_norm, False, AR_PUSH_MODE) + else: + if armed is not None: + missed( + ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden, residual = self._boundary(plan.attn_mc_ar, hidden, residual, post_norm) + + norm_w = _next_norm(core, i) + if i not in core.moe_layers: + # A dense layer: its MLP's down projection pushes into the post-MLP boundary's all-reduce, as the o_proj + # does into the attention one (the same slot; the gate_up projection never pushes). + dense_push = plan.attn_mc_ar and decode_graph and DENSE_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if dense_push else None + out, dense_pushed = self._swiglu_mlp(hidden, f"l{i}_mlp", True, armed, dense=True) + if armed is not None and dense_pushed: + taken(DENSE_PUSH.path, num_tokens) + return self._prepushed( + i, out, residual, norm_w, False, AR_PUSH_MODE, capture=capture + ) + if armed is not None: + missed( + DENSE_PUSH.path, + num_tokens, + "the dense MLP's down projection did not push its rows into the armed boundary all-reduce", + ) + return self._boundary(plan.attn_mc_ar, out, residual, norm_w, capture=capture) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + path = moe_path(num_tokens) + moe_push = None + if path == MOE_ENGINE.path: + # The engine runs inside the MoE; looked up (or found missing) here first so a push is armed only for an + # engine that will take it. + engine = core.engine(i, num_tokens, shared=True) + if ( + engine is not None + and plan.moe_mc_ar + and decode_graph + and MOE_PUSH.covers(num_tokens) + ): + moe_push = self._moe_push_targets(num_tokens) + elif path == MOE16_ENGINE.path and decode_graph: + # Router, routed and shared experts in one kernel; its bf16 sum goes to the boundary all-reduce. + engine = core.moe16_engine(i, num_tokens) + if engine is not None: + _record("m3_moe16_engine", num_tokens) + out = m3_moe16_engine(hidden.contiguous(), engine) + taken(MOE16_ENGINE.path, num_tokens) + return self._boundary( + plan.moe_mc_ar, out, residual, norm_w, quant=quant, capture=capture + ) + if path in (MOE16_ENGINE.path, MC_FINALIZE.path) and plan.moe_mc_ar: + if self._bucket_flags(num_tokens) is None: + self._no_bucket(MC_FINALIZE.path, num_tokens) + else: + (shared, fc2, scales, perm), _, _, _ = self._moe( + i, hidden, return_parts=False, do_finalize=False + ) + _record("m3_mc_finalize_allreduce_norm", num_tokens, quant) + fused = m3_mc_allreduce.m3_mc_finalize_allreduce_norm( + fc2, + perm.view(num_tokens, -1).to(torch.int32), + scales.reshape(num_tokens, -1), + shared, + residual, + norm_w, + core.rms_eps, + m3_mc_allreduce.get_workspace(core.mapping, core.hidden), + use_pdl=get_env_enable_pdl(), + quant=quant, + ) + if fused is None: + raise RuntimeError( + "the MC finalize all-reduce rejected the unfinalized MoE outputs" + ) + taken(MC_FINALIZE.path, num_tokens) + return fused + shared, routed, engine_pushed, shared_pushed = self._moe( + i, hidden, return_parts=True, push=moe_push + ) + if moe_push is not None: + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + if engine_pushed: + taken(MOE_PUSH.path, num_tokens, "engine, routed + shared") + return self._prepushed(i, shared, residual, norm_w, quant, 2, capture=capture) + else: + if engine_pushed and shared_pushed: + taken(MOE_PUSH.path, num_tokens, "routed, shared") + return self._prepushed( + i, shared, residual, norm_w, quant, 3, x2=routed, capture=capture + ) + if engine_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + missed( + MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary( + plan.moe_mc_ar, shared, residual, norm_w, x2=routed, quant=quant, capture=capture + ) + + def _swiglu_mlp( + self, + h: torch.Tensor, + prefix: str, + low_m_gemm: bool, + push: Optional[tuple] = None, + dense: bool = False, + ): + """A gated MXFP8 MLP (the dense MLP with ``dense``, the shared expert): gate_up, SwiGLU-OAI, down, no + reduction. Returns (output, whether the down projection's GEMV pushed into ``push``).""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm, push) + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: the MoE + scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor, do_finalize: bool): + """Router + routed experts: the router logits of a generation step (``_decode_router_logits``) and the experts + quantizing their bf16 input.""" + core = self.core + router_logits = _decode_router_logits(core, i, h) + if not do_finalize: + return _routed_experts(core, i, h, None, router_logits, False) + return self._routed_chunks(i, h, router_logits) + + def _moe( + self, i: int, h: torch.Tensor, *, return_parts: bool, do_finalize: bool = True, push=None + ): + """The MoE block of layer ``i`` without its reduction. + + ``push`` = (the engine's ``arp`` descriptor, the shared down projection's ``ar`` tuple) the decode layer armed + for the MoE boundary's exchange sets, or None. Returns ``(out, routed, engine_pushed, shared_pushed)``: with + ``return_parts`` the shared and routed parts (routed None when the engine folded the shared expert in), else + their sum in ``out``; with ``do_finalize=False`` ``out`` is ``[shared, fc2 out, expert scales, permuted + index]``. + """ + core = self.core + engine_arp, shared_ar = push if push is not None else (None, None) + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + do_finalize + and h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, engine_arp is not None) + folded = m3_moe_engine(h.contiguous(), engine, None, engine_arp) + taken(MOE_ENGINE.path, rows) + return folded, None, engine_arp is not None, False + routed_out, (shared_out, shared_pushed) = maybe_execute_in_parallel( + lambda: self._routed(i, h, do_finalize), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False, shared_ar), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the same + # collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if not do_finalize: + return [shared_out, *routed_out], None, False, shared_pushed + if return_parts: + return shared_out, routed_out, False, shared_pushed + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None, False, shared_pushed + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _generic_layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed, _, _ = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = None): + """Attention of layer ``i`` up to its o_proj (no reduction). ``push`` is the attention boundary's exchange the + o_proj may write into. Returns (o_proj output, pushed).""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v: an FP8 cache, or the FP8 pool the dense layers + # sit in within a hybrid NVFP4 cache on a decode step), the attention, o_proj. A covered (query rows, + # requests) shape attends on the full-context PrimTS decode (``FULLCTX_DECODE``), the norm/RoPE kernel then + # storing the new K / V rows itself (``DENSE_KV_INSERT_FUSED``); every other step the MSA backend. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + out_fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + plan = self._fullctx_plan(i, qkv.device, step) + kv_dest = None + if ( + plan is not None + and out_fp8 + and DENSE_KV_INSERT_FUSED + and DENSE_QKV_FP8_VIA_BF16 + and qkv.dtype == torch.bfloat16 + and core.head_dim == 128 + ): + kv_dest = ( + step["kv_cache_manager"].get_buffers(i, kv_layout="HND"), + step["msa_out_cache_loc"], + plan[1], + torch.empty( + (qkv.shape[0], core.heads, core.head_dim), + dtype=torch.float8_e4m3fn, + device=qkv.device, + ), + ) + fused = _qk_norm_rope(core, i, qkv, position_ids, out_fp8, kv_dest) + if kv_dest is not None: + taken(DENSE_KV_INSERT_PATH, kv_dest[2]) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + if kv_dest is not None: + q = kv_dest[3].view(q.shape[0], -1) + o = None + if plan is not None: + o = self._fullctx_decode(i, q, k, v, step, plan, kv_written=kv_dest is not None) + if o is None: + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _fullctx_plan(self, i: int, device: torch.device, step): + """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs on + kernels/m3_fullctx_decode: a covered (query rows per request, requests) shape over the layer's FP8 pool. None + otherwise (the MSA backend runs).""" + path = FULLCTX_DECODE.path + if is_torch_compiling(): + uncovered(path, "not a pure generation step on the MSA backend") + return None + span = step["msa_decode_span"] + if span is None or span.row_first != 0: + uncovered(path, "no decode span starting at row 0") + return None + nreq = int(step["num_seqs"]) + qlen = int(span.query_len) + if not FULLCTX_DECODE.covers((qlen, nreq)): + uncovered(path, f"{qlen} query rows x {nreq} requests is not a covered shape") + return None + if int(step["max_seq_len"]) > FULLCTX_DECODE_MAX_KV_LEN: + uncovered(path, "the cache length bound is beyond the plans' K/V capacity") + return None + num_tokens = nreq * qlen + pool, _ = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + if pool.dtype != torch.float8_e4m3fn: + missed(path, num_tokens, "the dense layer's KV pool is not FP8") + return None + own_reducer = FULLCTX_DECODE_OWN_REDUCER + if own_reducer and not _fullctx_decode_kernel.has_own_reducer(): + missed(path, num_tokens, "the BF16 split-KV merge op is not built") + return None + # Compiled before warmup (``prebuild_kernels``); the forward never compiles one. + plan = _fullctx_decode_kernel.plan_for(device, qlen, nreq, own_reducer) + if plan is None: + missed(path, num_tokens, "the PrimTS decode plan is not prebuilt") + return None + return plan, num_tokens + + def _fullctx_decode( + self, i: int, q, k, v, step, plan_tokens, *, kv_written: bool + ) -> torch.Tensor: + """Dense layer ``i``'s decode attention on the PrimTS plan of ``_fullctx_plan``: the step's new K / V go into + the paged cache as the MSA backend writes them (``write_msa_phase_kv``; already there when ``kv_written``), + then every query row attends its request's full causal context. The merge also emits the o_proj's MXFP8 + input (``FULLCTX_DECODE_MXFP8``) when the plan supports it.""" + core = self.core + backend = core.attn[i] + plan, num_tokens = plan_tokens + nreq = int(step["num_seqs"]) + if not kv_written: + write_msa_phase_kv( + backend, + k[:num_tokens], + v[:num_tokens], + step["metadata"], + AttentionInputType.mixed, + token_offset=0, + ) + pool, subpages = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + k_pages, v_pages = _fullctx_decode_kernel.kv_page_views(pool, subpages) + mxfp8 = None + if ( + FULLCTX_DECODE_MXFP8 + and output.shape[0] == num_tokens + and _fullctx_decode_kernel.emits_mxfp8() + and MERGE_MXFP8_MIN_TOKENS <= num_tokens <= MERGE_MXFP8.MAX_ROWS + ): + mxfp8 = MERGE_MXFP8.buffers(output.shape[1], output.device) + _record("m3_fullctx_decode", num_tokens, kv_written, mxfp8 is not None) + emitted = plan.run( + fp8_e4m3_saturate(q[:num_tokens].reshape(num_tokens, core.heads, core.head_dim)), + k_pages, + v_pages, + step["msa_block_table"][:nreq], + step["msa_seq_lens_cuda"][:nreq], + core.sm_scale(i), + output[:num_tokens], + mxfp8=mxfp8, + ) + if emitted: + m3_prequant.SLOT.publish(output, mxfp8[0][:num_tokens], mxfp8[1]) + taken(FULLCTX_DECODE_MXFP8_PATH, num_tokens) + taken(FULLCTX_DECODE.path, num_tokens) + return output + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step, push): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection and + its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and the FP8 + indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the o_proj. A + layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): a sparse + layer whose caches the producer wrote runs this target's selection and sparse decode when the generation span + starts at row 0; everything else the MSA backend (``m3_msa_attention``: the KV write unless a producer wrote the + caches, the sparse layers' top-k block selection, the FMHA).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + span = step["msa_decode_span"] + sparse = i in core.sparse_layers + if sparse and k_live is None and span is not None and span.row_first == 0: + self._decode_sparse_core(i, q_live, idx_q_live, step, out_live) + return output + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step of a sparse layer whose caches + the producer already wrote. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise the CuTe + DSL scorer writes the per-index-head block maxima, reduced to the KV head, and the long-row selector takes the + top 16 (the built-in selector beyond its shapes). + """ + core = self.core + backend = core.attn[i] + config = backend.m3_config + num_tokens = int(q.shape[0]) + rows = int(step["num_seqs"]) + qlen = int(step["msa_decode_span"].query_len) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + num_kv_heads = int(backend.indexer.config.num_kv_heads) + identity = step["msa_short_ok"] and step["msa_short_table"] is not None + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + else: + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = step["msa_idx_k_cache"](i) + max_score = step["msa_proxy_max_score_view"]( + config.num_index_heads, step["msa_worst_case_max_k_tiles"], num_tokens + ) + if INDEX_SCORE_T2SC_PATH.covers(rows): + # T2SC: this target's copy of the TP2 / TP8 scorer at the request count's split and prestage. + split_k, early = index_score_t2sc_schedule(rows) + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(split_k), + bool(early), + ) + if key not in _index_score_t2sc_kernel.IndexScoreRunner.kernel_cache: + missed( + INDEX_SCORE_T2SC_PATH.path, + num_tokens, + f"the T2SC index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score_t2sc", key) + m3_index_score_t2sc( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + split_k, + early, + ) + taken(INDEX_SCORE_T2SC_PATH.path, rows) + else: + early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(INDEX_SCORE_SPLIT_K), + bool(early), + ) + if key not in _index_score_kernel.IndexScoreRunner.kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + missed( + INDEX_SCORE.path, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score", key) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + INDEX_SCORE_SPLIT_K, + early, + ) + taken(INDEX_SCORE.path, rows) + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + group = config.num_index_heads // config.num_kv_heads + max_score_kv = ( + max_score.view( + config.num_kv_heads, group, max_score.shape[1], max_score.shape[2] + ).amax(dim=1) + if group > 1 + else max_score + ) + early_trigger = num_tokens <= SELECT_EARLY_MAX_ROWS + trigger_phase = ( + SELECT_LONG_TRIGGER_PHASE if num_tokens <= SELECT_LONG_TRIGGER_MAX_ROWS else 0 + ) + dsl = False + if LONG_SELECT_DSL.keys and 128 < max_score_kv.shape[1]: + # Compiled before warmup (prebuild_kernels); a missing kernel falls back to the 1024-thread one. + key = _select_long_dsl_kernel.kernel_key( + get_env_enable_pdl(), early_trigger, trigger_phase + ) + dsl = key in _select_long_dsl_kernel.kernel_cache + if not dsl: + missed( + LONG_SELECT_DSL.path, + num_tokens, + f"the 384-thread select {key} was not compiled before warmup", + ) + table = long_select_blocks( + max_score_kv, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + use_pdl=get_env_enable_pdl(), + early_trigger=early_trigger, + trigger_phase=trigger_phase, + fast_tail=SELECT_LONG_SL3, + dsl=dsl, + ) + if table is not None: + taken((LONG_SELECT_DSL if dsl else LONG_SELECT).path, num_tokens) + else: + uncovered( + LONG_SELECT.path, + f"{max_score_kv.shape[1]} KV blocks, outside {LONG_SELECT_MIN_BLOCKS}..{LONG_SELECT_MAX_BLOCKS}: " + "the built-in selector", + ) + table = m3_select_blocks( + max_score_kv, + MSA_REQUIRED_TOPK, + n_valid_blocks, + config.init_blocks, + config.local_blocks, + True, + ) + topk = table.permute(1, 0, 2) + manager = step["kv_cache_manager"] + k_paged, v_paged = msa_paged_kv(manager, i) + nvfp4 = {} + if manager.is_nvfp4_layer(i): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(backend, core.kv_scales[i]) + scales = manager.get_block_scale_buffers(i, "HND") + nvfp4 = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + _record("m3_sparse_decode", num_tokens, qlen, identity) + # The long-context decode (if it runs) prefetches this layer's o_proj weights into L2 and the blocks the next + # sparse layer selected at the previous step; any other route leaves the slots unused, and they are cleared + # after the call either way. + if SPARSE_DECODE_LC_MERGE_ONCE and SPARSE_DECODE_LC_PREFETCH_OPROJ: + m3_attn_lc.PREFETCH.tensors = (core.w[f"l{i}_o_w"], core.w[f"l{i}_o_s"]) + if SPARSE_DECODE_LC_MERGE_ONCE and i in core.n1_ids and num_tokens <= N1_MAX_ROWS: + m3_attn_lc.NEXT_KV.tensors = self._next_kv_prefetch(i, manager) + try: + m3_sparse_decode( + q.view(num_tokens, core.heads, core.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=core.sm_scale(i), + output=output.view(num_tokens, core.heads, core.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4, + ) + finally: + m3_attn_lc.PREFETCH.tensors = None + m3_attn_lc.NEXT_KV.tensors = None + + def _next_kv_prefetch(self, i: int, manager): + """(layer ``i``'s selected-block id buffer, the next sparse layer's, that layer's K / V pools and block scales) + for the long-context decode's next-layer prefetch; the last sparse layer's next is the first one (next step). + None when the next layer's cache is not NVFP4.""" + ids = self.core.n1_ids + order = sorted(ids) + nxt = next((j for j in order if j > i), order[0]) + if not manager.is_nvfp4_layer(nxt): + return None + k_next, v_next = msa_paged_kv(manager, nxt) + scales = manager.get_block_scale_buffers(nxt, "HND") + return (ids[i], ids[nxt], k_next, v_next, scales[:, 0], scales[:, 1]) + + +# -------------------------------------------------------------------------------------------------------------------- +# Speculative decoding: one-model Eagle3 at checkpoint precision +# -------------------------------------------------------------------------------------------------------------------- + + +class M3SpecStepFold: + """``SpecStepFold`` of this target. + + Between two decode graphs of a one-model Eagle3 step, the spec sampler stores the step's outputs into its per-slot + store and the engine gathers them back into its flat input buffers for the next step. This object does both in one + launch each (``m3_spec_tail.slot_scatter`` / ``step_gather``): integer copies only, the same values the engine's and + the sampler's own copies write. + + With ``in_forward`` (``STEP_SPEC_FOLD_IN_FORWARD``) the forward of a decode-only step also does both at its tail + (``store_in_forward``, one ``m3_fold_tail`` launch inside the step graph): it stores the step's outputs for the rows + ``begin_step`` armed and stages the next step's inputs for those rows in that order. The arming is a device-side row + count and slot list staged before the replay, so a captured graph whose step was not armed writes nothing. The + sampler then skips its store for exactly the armed rows (``sampler_rows_stored``), and the next step skips its + gathers when its rows are those rows (``begin_step``). + + One instance per target model; the target passes it to ``RuntimeHooks.spec_step_fold``. + """ + + def __init__( + self, in_forward: bool = False, store_copy_wait: bool = False, cpp_tail: bool = False + ) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + ] = None + self.in_forward = in_forward + # The forward's tail as this target's C++ op instead of the Triton kernel (m3_fold_tail.fold_tail(cpp=)). + self.cpp_tail = in_forward and cpp_tail + # Device arming read by the forward's tail: ctrl[0] = armed rows (0 = write nothing), slots[:rows]. + self._ctrl: Optional[torch.Tensor] = None + self._slots: Optional[torch.Tensor] = None + # The rows (sequence slots) the next forward stores and stages for, as begin_step armed them. + self._armed: Optional[tuple[int, ...]] = None + # The rows the previous forward staged the next step's inputs for (confirmed by sampler_rows_stored). + self._staged_for: Optional[tuple[int, ...]] = None + # Output buffers the tail ran on: the captured graphs' static outputs (by address) and the last eager step's. + self._graph_outputs: set[int] = set() + self._eager_outputs: Optional[torch.Tensor] = None + self._logged = False + # store_copy_wait: the sampler records this (external) event on its D2H side stream after copying the store + # to the host, and the forward's tail waits on it before it rewrites the store; the sampler can then copy + # on the side stream on the steps this forward stored the rows (see SpecSamplerBase._copy_outputs_to_host). + self._store_copy_wait = in_forward and store_copy_wait + self.store_copies_done: Optional[torch.cuda.Event] = None + + def bind_engine_buffers( + self, + input_ids: torch.Tensor, + draft_tokens: torch.Tensor, + position_id_offsets: torch.Tensor, + kv_lens_offsets: torch.Tensor, + max_batch_size: int, + ) -> None: + self.step_inputs = (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) + if self._store_copy_wait: + # Recorded once now so a captured wait has a completed record to wait on before the first copy. + self.store_copies_done = torch.cuda.Event(external=True) + self.store_copies_done.record() + logger.info( + "[M3 host] spec store host copies on the D2H side stream, awaited by the forward's tail" + ) + if self.cpp_tail: + logger.info( + "[M3 host] spec fold tail: one C++ launch (no Triton) at the forward's tail" + ) + if self.in_forward: + self._ctrl = torch.zeros(1, dtype=torch.int32, device=input_ids.device) + self._slots = torch.zeros( + max(max_batch_size, 1), dtype=torch.long, device=input_ids.device + ) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: + if not self.in_forward or self._ctrl is None: + return False + rows = tuple(seq_slots) if seq_slots is not None else None + staged = rows is not None and rows == self._staged_for + self._staged_for = None + armed = rows if rows is not None and 0 < len(rows) <= self._slots.numel() else None + # Read only by this step's forward, so both uploads may wait for the step-input pack's flush. + if armed is not None: + slots = torch.tensor(armed, dtype=torch.long, pin_memory=prefer_pinned()) + copy_h2d_if_changed(self._slots[: len(armed)], slots, pack=True) + ctrl = torch.tensor( + [len(armed) if armed is not None else 0], dtype=torch.int32, pin_memory=prefer_pinned() + ) + copy_h2d_if_changed(self._ctrl, ctrl, pack=True) + self._armed = armed + return staged + + # The launches the store wait may precede. store_copies_done is an external event: under graph capture its wait + # is a node of its own, so the kernel launched right after it loses its programmatic (PDL) edge and with it the + # overlap of its pre-wait part. Only a launch that waits for its predecessor's completion anyway may follow it: + # the fold tail (a plain launch, no griddepcontrol). + STORE_WAIT_BEFORE = frozenset({"fold_tail"}) + + def _wait_store_copies(self, before: str) -> None: + """Order the store's rewrite after the sampler's host copies of it; ``before`` names the launch that follows. + Under graph capture a launch outside STORE_WAIT_BEFORE fails the capture.""" + if self.store_copies_done is None: + return + if before not in self.STORE_WAIT_BEFORE and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"M3SpecStepFold: the store wait would precede {before!r}, a launch outside " + f"{sorted(self.STORE_WAIT_BEFORE)}; an external-event wait node drops the next kernel's PDL edge" + ) + torch.cuda.current_stream().wait_event(self.store_copies_done) + + def store_in_forward(self, outputs: dict) -> None: + """At the forward's tail: store the armed rows' outputs and stage the next step's inputs (one launch).""" + if self._ctrl is None or self.store is None or self.step_inputs is None: + return + store = self.store + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + new_tokens = outputs.get("new_tokens") + next_new_tokens = outputs.get("next_new_tokens") + new_tokens_lens = outputs.get("new_tokens_lens") + next_draft_tokens = outputs.get("next_draft_tokens") + if any( + t is None or not isinstance(t, torch.Tensor) + for t in (new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens) + ): + return + rows = new_tokens.shape[0] + # The previous step's host copies of the store run on the sampler's side stream. + self._wait_store_copies(before="fold_tail") + if not m3_fold_tail.fold_tail( + self._ctrl, + self._slots, + min(rows, self._slots.numel()), + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + next_draft_tokens, + store.next_draft_tokens, + new_tokens_lens, + store.new_tokens_lens, + input_ids, + draft_tokens, + position_id_offsets, + kv_lens_offsets, + next_new_tokens.shape[1], + cpp=self.cpp_tail, + ): + return + if torch.cuda.is_current_stream_capturing(): + self._graph_outputs.add(new_tokens.data_ptr()) + else: + self._eager_outputs = new_tokens + + def gather( + self, + new_tokens: torch.Tensor, + next_draft_tokens: torch.Tensor, + new_tokens_lens: torch.Tensor, + previous_slots: torch.Tensor, + previous_pos_indices: torch.Tensor, + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + if self.step_inputs is None: + return False + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + return m3_spec_tail.step_gather( + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len, + tokens_per_gen_step, + draft_token_buffer_width, + input_ids, + num_tokens, + draft_tokens, + num_draft_tokens, + position_id_offsets, + kv_lens_offsets, + first_previous_row, + ) + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list[int]) -> bool: + armed, self._armed = self._armed, None + new_tokens = outputs.get("new_tokens") + ran = isinstance(new_tokens, torch.Tensor) and ( + new_tokens is self._eager_outputs or new_tokens.data_ptr() in self._graph_outputs + ) + self._eager_outputs = None + stored = ( + ran + and armed is not None + and num_skip == 0 + and tuple(seq_slots) == armed + and new_tokens.shape[0] == len(armed) + ) + self._staged_for = armed if stored else None + if stored and not self._logged: + self._logged = True + logger.info("[M3 host] spec step stored and next step staged in the forward tail") + return stored + + def scatter( + self, + seq_slots: torch.Tensor, + new_tokens: torch.Tensor, + next_new_tokens: torch.Tensor, + new_tokens_lens: torch.Tensor, + next_draft_tokens: torch.Tensor, + ) -> bool: + store = self.store + if store is None: + return False + return m3_spec_tail.slot_scatter( + seq_slots, + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + new_tokens_lens, + store.new_tokens_lens, + next_draft_tokens, + store.next_draft_tokens, + ) + + +class _DraftForward: + """``draft_model.model`` for the in-tree worker: the drafter layer's forward, called as + ``model(input_ids=, position_ids=, hidden_states=, attn_metadata=, spec_metadata=, all_rank_num_tokens=)`` and + returning (hidden, hidden to carry to the next pass). ``d2t`` None: the draft vocabulary is the target's.""" + + d2t = None + + def __init__(self, draft: "M3Eagle3Draft") -> None: + self._draft = draft + + def __call__(self, **kwargs): + return self._draft.forward_layer(**kwargs) + + +class M3Eagle3Draft(nn.Module): + """The one-layer BF16 Eagle3 drafter (``speculative_model``): its weights (``weights.DRAFT_WEIGHTS``), its + vocab-parallel BF16 head and logits processor, the drafter layer's TRTLLM attention backend (layer index = the + target's layer count, FP8 KV) and its TP all-reduce. + + What the worker reads: ``model`` (the layer forward), ``apply_eagle3_fc`` (the per-layer fc_norms + fc over the + captured target hidden states), ``lm_head`` and ``logits_processor``; the fused draft loop reads the layer's parts + (``draft_input_weights``, ``attention``, ``mlp``, ``ar_norm``). The drafter embeds with the target's table. + """ + + def __init__( + self, + core: MiniMaxM3Tp4ModelingV2, + draft_config: ModelConfig, + start_layer_idx: int, + capture_layers: int, + ) -> None: + super().__init__() + cfg = draft_config.pretrained_config + _ensure_draft_vocab_size(cfg) + eagle = getattr(cfg, "eagle_config", None) or {} + mapping = draft_config.mapping + self.draft_config = draft_config + # The engine's draft weight mapper binds to these two. + self.model_config = draft_config + self.config = cfg + self.capture_layers = int(capture_layers) + self.layer_idx = int(start_layer_idx) + self.dtype = cfg.torch_dtype + self.hidden = int(cfg.hidden_size) + self.hidden_size_in = int(getattr(cfg, "target_hidden_size", cfg.hidden_size)) + self.rms_eps = float(cfg.rms_norm_eps) + self.return_hidden_post_norm = bool( + eagle.get("return_hidden_post_norm", False) or getattr(cfg, "norm_output", False) + ) + if ( + int(cfg.num_hidden_layers) != 1 + or eagle.get("eh_proj_before_attn", False) + or eagle.get("norm_before_fc", False) + or not getattr(cfg, "fc_norm", False) + or getattr(cfg, "kv_lora_rank", None) + or self.capture_layers <= 1 + or self.hidden_size_in != self.hidden + or cfg.draft_vocab_size != cfg.vocab_size + or getattr(cfg, "attention_bias", False) + or getattr(cfg, "mlp_bias", False) + or getattr(cfg, "bias", False) + or getattr(cfg, "tie_word_embeddings", False) + ): + raise ValueError( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle3 drafter with per-layer " + "fc_norms over " + "several captured layers, the target's hidden size and vocabulary and no biases" + ) + tp = mapping.tp_size + self.heads = cfg.num_attention_heads // tp + self.kv_heads = (cfg.num_key_value_heads + tp - 1) // tp + self.head_dim = int( + getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + ) + self.q_size = self.heads * self.head_dim + self.kv_size = self.kv_heads * self.head_dim + self.w = _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + self.lm_head = LMHead( + cfg.draft_vocab_size, + cfg.hidden_size, + dtype=cfg.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=draft_config.lm_head_gather_output, + reduce_output=False, + use_custom_cublas_mm=False, + quant_config=None, + ) + self.logits_processor = LogitsProcessor() + self.allreduce = AllReduce( + mapping=mapping, strategy=draft_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + quant_config = draft_config.get_quant_config() + self.attn = create_attention( + draft_config.attn_backend, + self.layer_idx, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, rope=RopeParams.from_config(cfg) + ), + quant_config=quant_config, + skip_create_weights_in_init=draft_config.skip_create_weights_in_init, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): + raise ValueError( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + # The drafter embeds with the target's table, read off the core on every call (the engine materializes the + # core's weights after construction). + self._core = [core] + self.model = _DraftForward(self) + # The projections (weight keys) on the fixed-tactic low-M kernels (``DRAFT_LOWM``), set by + # ``prebuild_linears``. + self._lowm_keys: frozenset = frozenset() + + def prebuild_linears(self) -> int: + """Put the drafter's BF16 projections ``DRAFT_LOWM`` covers on the fixed-tactic low-M kernels and compile + every covered call. Once the weights are real, before the executor's warmup, on every rank. Returns the + number of low-M calls compiled.""" + keys = ("fc", "qkv", "o", "gu", "dn") + built = _draft_lowm_kernel.prebuild([self.w[k] for k in keys], DRAFT_LOWM.keys) + self._lowm_keys = frozenset(k for k in keys if _draft_lowm_kernel.installed(self.w[k])) + return built + + def _linear(self, x: torch.Tensor, key: str) -> torch.Tensor: + """``x @ w[key]^T``: the low-M kernel at the rows ``DRAFT_LOWM`` covers (``prebuild_linears``), else the plain + GEMM.""" + w = self.w[key] + if key in self._lowm_keys: + out = _draft_lowm_kernel.draft_linear(w, x, None) + if out is not None: + return out + return m3_bf16_linear(x, w) + + @property + def embed_table(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def load_weights_from_target_model(self, target_model: nn.Module) -> None: + """Nothing to alias: the drafter reads the target's embedding table through the core and owns its head.""" + + # ------------------------------------------------------------------------------------------------- the layer + + def fc(self, hs: torch.Tensor) -> torch.Tensor: + return self._linear(hs, "fc") + + def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The captured target hidden states [T, layers x H] to the drafter's input width: each layer's fc_norm, + then the fc (the in-tree ``Eagle3ForCausalLM.apply_eagle3_fc``).""" + hidden_states = hidden_states.to(self.dtype) + if hidden_states.shape[-1] != self.hidden: + chunks = hidden_states.chunk(self.capture_layers, dim=-1) + hidden_states = torch.cat( + [ + flashinfer_rmsnorm(c, self.w[f"fc_norm{i}"], self.rms_eps) + for i, c in enumerate(chunks) + ], + dim=-1, + ) + hidden_states = self.fc(hidden_states) + return hidden_states + + def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: + """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is + False, where the caller fuses the reduction).""" + qkv = self._linear(x, "qkv") + _record("m3_draft_attention", int(md.num_tokens)) + o = m3_draft_attention(self.attn, qkv, md) + out = self._linear(o, "o") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool, silu_cpp: bool = False) -> torch.Tensor: + """gate_up, the plain SwiGLU (with ``silu_cpp``, ``DRAFT_SILU_CPP``: this target's C++ op, bit for bit the + Triton one, in one launch), down.""" + gu = self._linear(h, "gu") + act = None + if silu_cpp and DRAFT_SILU_CPP: + act = m3_draft_silu(gu) + if act is None: + uncovered( + DRAFT_SILU_PATH, "the drafter SwiGLU input layout (the Triton op runs it)" + ) + else: + taken(DRAFT_SILU_PATH, int(gu.shape[0]), "drafter MLP SwiGLU, one C++ launch") + if act is None: + act = m3_silu_and_mul(gu, None, None, None) + out = self._linear(act, "dn") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def ar_norm(self, x: torch.Tensor, residual: torch.Tensor, norm_key: str): + """One draft-layer boundary in the in-tree structure: the projection's plain all-reduce (BF16 out), then the + fused residual add + RMSNorm (FP32 add, norm of the FP32 sum). Returns (normed, residual sum).""" + return flashinfer_fused_add_rmsnorm( + m3_allreduce(x, self.allreduce, None), residual, self.w[norm_key], self.rms_eps + ) + + def forward_layer( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + all_rank_num_tokens: Optional[List[int]] = None, + ): + """The drafter layer (the in-tree ``Eagle3DraftModel.forward`` for one layer): the embedded tokens and the + fc'd hidden states normed and concatenated, attention, the post-attention norm with the residual, the MLP and + the final norm. Returns (normed output, the hidden state the next pass carries).""" + previous = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + embeds = m3_embedding(input_ids, self.embed_table).to(self.dtype) + residual = hidden_states + h = flashinfer_rmsnorm(hidden_states, self.w["hidden_norm"], self.rms_eps) + e = flashinfer_rmsnorm(embeds, self.w["in_norm"], self.rms_eps) + h = torch.cat([e, h], dim=-1) + h = self.attention(h, attn_metadata, reduce=True) + h, residual = flashinfer_fused_add_rmsnorm( + h, residual, self.w["post_norm"], self.rms_eps + ) + h = self.mlp(h, reduce=True) + spec_metadata.maybe_capture_hidden_states(self.layer_idx, h, residual) + h, to_save = flashinfer_fused_add_rmsnorm( + h, residual, self.w["final_norm"], self.rms_eps + ) + return (h, h) if self.return_hidden_post_norm else (h, to_save) + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous + + +def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, capture_norm) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "tp_size", 1) > 1 + ) + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify and fused draft loop, at the in-tree numerics. + + * Verify: on decode-only batches whose verify rows the fused kernels cover (``tail_eligible``), the target's BF16 + lm_head computes this rank's vocabulary shard of the logits (the in-tree head's own local GEMM; at most + ``TARGET_HEAD_MAX_ROWS`` verify rows stream a lossless 12-bit coding of the same BF16 weight) and one kernel + replaces the logits all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. + * Draft: the drafter's first pass may take the per-layer fc_norms in one launch (or folded into the hidden-state + capture) and run its fc beside the target head GEMM (forked right after the last captured target layer at most + ``FC_EARLY_MAX_ROWS`` verify rows). On a decode-only step ``DRAFT_STEP`` covers, the draft loop is fused: each + pass's draft input in one launch, the drafter layer's attention and MLP with their all-reduces fused with the + residual add and the next RMSNorm, and one kernel for the BF16 head, the cross-rank argmax and the pass's + bookkeeping. + + Every other batch takes the in-tree paths. Methods whose docstring says "Mirrors" copy the named in-tree method + with the fused branches inserted. + """ + + def __init__(self, spec_config, mapping, use_separate_draft_kv_cache: bool = False) -> None: + super().__init__( + spec_config, mapping=mapping, use_separate_draft_kv_cache=use_separate_draft_kv_cache + ) + # Per-step state, reset at the top of every _forward_impl: the first draft pass's gather rows when this step + # takes the fused first drafter inputs, and the fc output a fork computed beside the target head. + self._m3_first_gather_ids: Optional[torch.Tensor] = None + self._m3_fc_hidden: Optional[torch.Tensor] = None + self._m3_target_lm_head: Optional[nn.Module] = None + # (fc_norm0 weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple] = None + # Per step: the captured layers written with the capture norm (set by capture_hidden_states). + self._m3_capture_normed: set = set() + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + # Armed per step by arm_capture_norm: the drafter whose first fc is forked at the last captured target layer, + # and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + # The 12-bit coded copy of the target head shard the fused verify streams, built before the warmup. + self._m3_target_head = None + self._m3_force_vec_warm = False + # The streamed copy of the drafter's BF16 head shard the fused draft loop reads, built before the warmup. + self._m3_draft_head = None + # Per step: the next step's input tokens [nreq, D + 1], column 0 written by the fused verify (or the seed) and + # columns 1.. by the fused draft passes; set as _m3_next_new_tokens once the fused draft loop filled it. + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + # Per step: the first draft inputs are the fused ones (accepted tokens as they lie, fc output of every row). + self._m3_fused_first_inputs = False + + def capture_target( + self, spec_metadata, layer_id: int, num_tokens: int + ) -> Optional[Tuple[torch.Tensor, torch.Tensor, float]]: + """Where this step's capture norm writes layer ``layer_id``'s capture, for a boundary all-reduce that writes it + itself (``CAPTURE_NORM_IN_BOUNDARY``): (the capture buffer's [num_tokens, H] column block, the layer's + fc_norm weight, eps), or None (no capture norm armed, not a captured layer, or more rows than the buffer). + The boundary writes every row it reduces; rows past num_capture_tokens are unread.""" + capture_norm = self._m3_capture_norm + if capture_norm is None or spec_metadata.hidden_states is None: + return None + if num_tokens > spec_metadata.hidden_states.shape[0]: + return None + hidden_size = spec_metadata.hidden_size + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + out = spec_metadata.hidden_states[ + :num_tokens, i * hidden_size : (i + 1) * hidden_size + ] + return out, capture_norm[0][i], capture_norm[1] + return None + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None, written: bool = False + ) -> None: + """The Eagle3 capture with this step's capture norm (``written``: the layer's boundary all-reduce already + wrote it, see ``capture_target``); on a step armed for the early fc fork (``arm_capture_norm``), the last + captured layer then forks the drafter's first fc onto the aux stream, where it overlaps the remaining target + layers and the head (joined by ``_m3_target_accept``).""" + if not written: + _eagle3_capture(spec_metadata, layer_id, hidden_states, residual, self._m3_capture_norm) + if self._m3_capture_norm is not None: + # Written with this step's capture norm: by the boundary all-reduce (written) or just above. + self._m3_capture_normed.add(layer_id) + drafter = self._m3_fc_early_drafter + layers = spec_metadata.layers_to_capture + if drafter is not None and layers and layer_id == layers[-1]: + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = self._m3_fc_fork( + spec_metadata, spec_metadata.num_capture_tokens, drafter + ) + + def prebuild(self, lm_head: nn.Module, draft: Optional[M3Eagle3Draft]) -> None: + """The fused verify's coded target head and the fused draft loop's head copy, from the loaded weights, + before the executor's warmup, on every rank.""" + start = time.perf_counter() + if self.prebuild_target_head(lm_head): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] coded target head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + start = time.perf_counter() + if self.prebuild_draft_head(draft): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Mirrors Eagle3OneModelWorker._forward_impl, with the fused verify branch.""" + runtime_draft_len = spec_metadata.runtime_draft_len + # skip the draft forward if the runtime draft length is 0 + if runtime_draft_len == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._m3_first_gather_ids = None + self._m3_next_new = None + self._m3_next_new_tokens = None + self._m3_fused_first_inputs = False + if logits is None: + # tail_eligible() skipped the gathered target logits: verify from this rank's logits shard. + accepted_tokens, num_accepted_tokens, self._m3_first_gather_ids = ( + self._m3_target_accept( + hidden_states, + spec_metadata, + batch_size, + runtime_draft_len, + draft_model, + kv_lens=getattr(attn_metadata, "kv_lens_cuda", None), + ) + ) + else: + self._execute_guided_decoder_if_present(logits) + + # Sample and accept tokens. ``input_ids`` is required by the relaxed-acceptance path (scans for + # thinking-phase tokens); ignored otherwise. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, attn_metadata, spec_metadata + ) + if self.draft_only_eligible(attn_metadata, spec_metadata, draft_model): + self._m3_seed_first_pass(num_accepted_tokens, spec_metadata, batch_size) + + # Mamba hybrid models need state updates after token acceptance (a no-op on this target's cache manager). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None) + if self.sa_enhancer is not None and sa_manager is not None: + self.sa_enhancer.extend_and_prepare( + sa_manager=sa_manager, + request_ids=spec_metadata.request_ids, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_gens=num_gens, + num_contexts=num_contexts, + max_draft_len=runtime_draft_len, + ) + + if self._uses_external_shared_target_kv: + next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + ) + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here so the post-loop restore (below) + # can put attn_metadata back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + # prepare next new tokens to support overlap scheduler + if self._m3_next_new_tokens is not None: + # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). + next_new_tokens = self._m3_next_new_tokens + else: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + # Commit model-owned target state only after draft preparation has succeeded; SpecWorkerBase aborts the + # still-pending snapshots on any exception above this point. + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices[:batch_size], + num_contexts, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """The fused first draft inputs on a step that verified with the fused kernels (or was seeded for them), + else the in-tree ones.""" + fused = self._m3_first_drafter_inputs( + position_ids, + accepted_tokens, + input_ids.shape[0], + attn_metadata.num_contexts, + spec_metadata, + draft_model, + ) + self._m3_fused_first_inputs = fused is not None + if fused is not None: + taken("m3_first_drafter_inputs", attn_metadata.num_seqs) + fused.update(attn_metadata=attn_metadata, spec_metadata=spec_metadata) + return fused + return super().prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + def _m3_structural_ok(self, attn_metadata, spec_metadata) -> bool: + """The batch conditions shared by the fused verify and the fused first drafter inputs: a decode-only plain + Eagle3 step without attention DP, guided decoding, relaxed acceptance or a suffix-automaton enhancer, of at + most DECODE_MAX_TOKENS verify rows. Host state only, so the TP ranks agree.""" + if self.is_mtp_eagle or self.use_dynamic_tree: + return False + if self._uses_external_shared_target_kv or self.sa_enhancer is not None: + return False + if ( + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp + ): + return False + if getattr(self.spec_config, "use_relaxed_acceptance_for_thinking", False): + return False + if attn_metadata.num_contexts != 0 or attn_metadata.num_seqs == 0: + return False + draft_len = spec_metadata.runtime_draft_len + if draft_len == 0 or attn_metadata.num_seqs * (draft_len + 1) > DECODE_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Whether this step verifies with the fused kernels instead of the gathered target logits. + + Decode-only Eagle3 batches (``_m3_structural_ok``) over an unpadded vocab-sharded bias-free target head, all + greedy at a request count ``SHARD_ACCEPT`` covers, or sampled under the temperature + top-p contract (see + _m3_sampled_verify_ok), with nothing else reading the target logits. Every input is host state shared by all + TP ranks, so the ranks agree. + """ + if not self._m3_structural_ok(attn_metadata, spec_metadata): + return False + if getattr(spec_metadata, "enable_penalty", False): + return False + if not _sharded_head(lm_head) or getattr(lm_head, "bias", None) is not None: + return False + draft_len = spec_metadata.runtime_draft_len + if spec_metadata.is_all_greedy_sample: + key = (draft_len, attn_metadata.num_seqs) + if not SHARD_ACCEPT.covers(key): + uncovered( + SHARD_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + elif attn_metadata.num_seqs > SAMPLED_VERIFY_MAX_REQUESTS: + uncovered("m3_sampled_accept", f"sampled verify at {attn_metadata.num_seqs} requests") + return False + elif not self._m3_sampled_verify_ok(spec_metadata): + return False + self._m3_target_lm_head = lm_head + return True + + def draft_only_eligible(self, attn_metadata, spec_metadata, draft_model) -> bool: + """Whether a step verified by the in-tree path can still take the fused first drafter inputs: the structural + conditions of tail_eligible, without the target-sampling ones, up to ``FIRST_DRAFTER_INPUTS_MAX_REQUESTS`` + requests.""" + if attn_metadata.num_seqs > FIRST_DRAFTER_INPUTS_MAX_REQUESTS: + uncovered( + "m3_first_drafter_inputs", f"in-tree verify at {attn_metadata.num_seqs} requests" + ) + return False + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None + + def _m3_seed_first_pass(self, num_accepted_tokens, spec_metadata, batch_size: int) -> None: + """The first draft pass's gather rows (the last accepted row of each request) from the in-tree acceptance, as + _m3_target_accept leaves them.""" + k1 = spec_metadata.runtime_draft_len + 1 + n = num_accepted_tokens[:batch_size].to(torch.int64) + rows = spec_metadata.batch_indices_cuda[:batch_size].to(torch.int64) + self._m3_first_gather_ids = rows * k1 + n - 1 + + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], + ): + """The in-tree first draft inputs of a decode-only step, with the per-layer fc_norm RMSNorms of + ``apply_eagle3_fc`` in one kernel (unless the capture already applied them) before the fc; the input ids are + the accepted tokens as they lie. None to use the in-tree path.""" + if ( + self._m3_first_gather_ids is None + or self.is_mtp_eagle + or num_contexts != 0 + or draft_model is None + ): + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + ): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not draft_model.w["fc_norm0"]: + if torch.cuda.is_current_stream_capturing(): + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), + ) + self._m3_fc_norm_w = cache + hidden = self._m3_fc_hidden + self._m3_fc_hidden = None + if hidden is None or hidden.shape[0] != num_tokens: + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + hs = m3_spec_tail.chunk_norm(hs, cache[1], draft_model.rms_eps) + else: + hs = self._m3_renorm_raw_captures(hs, spec_metadata, draft_model) + hidden = draft_model.fc(hs) + input_ids = accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].reshape(-1) + return {"input_ids": input_ids, "position_ids": position_ids, "hidden_states": hidden} + + def _m3_renorm_raw_captures( + self, hs: torch.Tensor, spec_metadata, draft_model: M3Eagle3Draft + ) -> torch.Tensor: + """The captured rows ``hs`` [rows, layers * H] for a reader that skips the fc_norm (the capture norm is armed): + unchanged when every captured layer was written with this step's norm (capture_hidden_states records them), + else -- not expected on any path -- a copy with each layer captured raw normalized by its own fc_norm here, + logged as a miss, so a raw row never reaches the drafter's fc.""" + missing = [ + i + for i, layer_id in enumerate(spec_metadata.layers_to_capture) + if layer_id not in self._m3_capture_normed + ] + if not missing: + return hs + missed( + "m3_capture_norm", + int(hs.shape[0]), + f"captured layer(s) {[spec_metadata.layers_to_capture[i] for i in missing]} were written without the " + "capture norm; their fc_norm is applied before the drafter's fc", + ) + hs = hs.clone() + hidden = spec_metadata.hidden_size + for i in missing: + block = hs[:, i * hidden : (i + 1) * hidden] + block.copy_( + flashinfer_rmsnorm( + block.contiguous(), draft_model.w[f"fc_norm{i}"], draft_model.rms_eps + ) + ) + return hs + + def _m3_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + after: Optional[torch.cuda.Event] = None, + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (after the event + ``after`` when given, else after everything queued on the current stream); the result is picked up by + _m3_first_drafter_inputs. Returns the stream to join, or None when the fc path does not apply (the first + drafter inputs then compute it).""" + self._m3_fc_hidden = None + if draft_model is None: + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + or self._m3_fc_norm_w is None + ): + return None + aux = self._m3_aux_stream + if aux is None: + aux = self._m3_aux_stream = torch.cuda.Stream() + if after is not None: + aux.wait_event(after) + else: + aux.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(aux): + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + hs = m3_spec_tail.chunk_norm( + hs, self._m3_fc_norm_w[1], draft_model.rms_eps, use_pdl=False + ) + else: + hs = self._m3_renorm_raw_captures(hs, spec_metadata, draft_model) + self._m3_fc_hidden = draft_model.fc(hs) + return aux + + def _m3_fc_join(self, fork: torch.cuda.Stream) -> None: + """Join the fc fork onto the current stream and keep its output's memory alive for that stream. + + _m3_fc_hidden is allocated in fork (aux) stream order and read on the current stream. Without the record, + dropping it returns its block to the aux stream's pool, where later aux work can take it while the current + stream's reads are still queued (eager warm-up / mixed steps). Nothing is recorded under graph capture, as in + maybe_execute_in_parallel: a captured graph relies on the next fork's aux.wait_stream(main) to order the + reuse.""" + main = torch.cuda.current_stream() + main.wait_stream(fork) + hidden = self._m3_fc_hidden + if ( + hidden is not None + and hidden.is_cuda + and not (torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()) + ): + hidden.record_stream(main) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model: Optional[M3Eagle3Draft] + ) -> None: + """Before every target forward: whether this step's Eagle3 captures apply the drafter's fc_norm themselves. + + On a fused-verify step whose fc_norm weights are already stacked (by an earlier fused first pass), the + capture writes rmsnorm(h + r) * fc_norm[i] (one kernel per captured layer instead of add + copy) and the fused + first drafter inputs read it without their own chunk_norm. A fused-verify step of at most + ``FC_EARLY_MAX_ROWS`` verify rows is also armed to fork the drafter's first fc at the last captured layer + (``capture_hidden_states``). + """ + self._m3_capture_norm = None + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = None + # The captured layers this step wrote with its capture norm (capture_hidden_states). + self._m3_capture_normed = set() + # A fork result no first draft pass consumed must not reach a later step's (same-size) first inputs. + self._m3_fc_hidden = None + if draft_model is None: + return + if not self.tail_eligible(attn_metadata, spec_metadata, lm_head): + return + # The fork reads every captured column at the last captured layer, so a final-state capture, which is written + # after the model returns, keeps the fc at its stock place. + if ( + attn_metadata.num_seqs * (spec_metadata.runtime_draft_len + 1) <= FC_EARLY_MAX_ROWS + and FINAL_HIDDEN_LAYER_ID not in spec_metadata.layers_to_capture + ): + self._m3_fc_early_drafter = draft_model + fc_norm_w = self._m3_fc_norm_w + if ( + fc_norm_w is not None + and fc_norm_w[0] is draft_model.w["fc_norm0"] + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.rms_eps) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata) -> bool: + """Whether sampled verify rows can use the fused sampled op: only under advanced_sampling_mode=no_topk, the + deploy-time guarantee that no row uses top-k or min-p (that op applies neither).""" + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + return mode is not None and mode.skips_top_k and not mode.skips_top_p + + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + kv_lens: Optional[torch.Tensor] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path. ``kv_lens``: the + requests' KV lengths at this verify (the per-request forced-acceptance key, ``FORCE_AL_PER_REQUEST``).""" + force_total = 0 + force_vec = None + if self.force_num_accepted_tokens != 0.0: + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # The draw pool is built by an H2D copy and the draw kernel JIT-compiles per block width, so both + # happen here, before the first captured step (a scratch counter keeps the real draw sequence). + self._ensure_force_accept_rng_state(hidden_states.device) + if not self._m3_force_vec_warm: + scratch = torch.zeros(1, dtype=torch.int64, device=hidden_states.device) + dummy = torch.empty(64, dtype=torch.int32, device=hidden_states.device) + for block in (16, 32, 64): + m3_force_vec( + self._force_accept_rng_pool, + scratch, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + m3_force_vec_req( + self._force_accept_rng_pool, + dummy, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + self._m3_force_vec_warm = True + if not is_warmup: + force_total = min(int(self.force_num_accepted_tokens) + 1, draft_len + 1) + force_vec = self._m3_force_vec( + batch_size, force_total, draft_len + 1, hidden_states.device, kv_lens + ) + draft_tokens = spec_metadata.draft_tokens.view(batch_size, draft_len) + if draft_tokens.dtype != torch.int32: + draft_tokens = draft_tokens.to(torch.int32) + # Column 0 receives the last accepted token, the next step's first input (the fused draft loop adds the rest). + self._m3_next_new = torch.empty( + (batch_size, draft_len + 1), dtype=torch.int32, device=hidden_states.device + ) + rows = hidden_states[: batch_size * (draft_len + 1)] + lm_head = self._m3_target_lm_head + batch_idx = spec_metadata.batch_indices_cuda[:batch_size] + head = self._m3_coded_head(lm_head, rows.shape[0]) + # The Eagle3 fc on the aux stream: the fork made at the last captured layer, else one beside the head. A + # sampled verify's coded head holds every SM while it streams: its fc is forked after it instead (overlapping + # the sampled acceptance), where it does not delay the head's start. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + head_first = ( + fork is None + and (head is not None or rows.shape[0] <= FC_AFTER_HEAD_MAX_ROWS) + and not spec_metadata.is_all_greedy_sample + ) + if fork is None and not head_first: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the head's own local GEMM (or its coded copy), whose + # all-gather the fused kernels replace (they exchange per-row records across ranks, not logits). + if head is not None: + taken(TARGET_HEAD.path, rows.shape[0]) + logits = m3_target_head_kernel.logits(head, rows) + else: + logits = lm_head.apply_linear(rows, None) + head_done = None + if head_first: + head_done = torch.cuda.Event() + head_done.record() + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept_kernel.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = SHARD_ACCEPT.path + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept_kernel.sampled_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + spec_metadata.temperatures[:n_rows], + spec_metadata.top_ks[:n_rows], + spec_metadata.top_ps[:n_rows], + None, + seeds, + offsets, + force_total, + self._m3_next_new, + force_vec=force_vec, + one_round_k=SAMPLED_ACCEPT_1R_K if SAMPLED_ACCEPT_SM107TP4 else 0, + prewait_params=SAMPLED_ACCEPT_PREWAIT_PARAMS, + early_trigger=SAMPLED_ACCEPT_EARLY_TRIGGER, + ilp_walk=SAMPLED_ACCEPT_ILP_WALK, + ) + path = "spec_tail" + if head_done is not None: + # Enqueued after the sampled acceptance (which starts first) but waits on the head only. + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model, after=head_done) + if fork is not None: + self._m3_fc_join(fork) + if result is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, None, "the fused verify's exchange workspace") + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + taken( + SHARD_ACCEPT.path if spec_metadata.is_all_greedy_sample else "m3_sampled_accept", + batch_size, + f"draft length {draft_len}" + + ( + "" + if spec_metadata.is_all_greedy_sample + else f", SA-1R one_round_k {SAMPLED_ACCEPT_1R_K if SAMPLED_ACCEPT_SM107TP4 else 0}" + ), + ) + return result + + def prebuild_target_head(self, lm_head: nn.Module) -> bool: + """Build the 12-bit coded copy of the target head shard the fused verify streams at <= TARGET_HEAD_MAX_ROWS + rows (rebuilt from the current weight values on every call), on every rank. True when the coded copy decodes + to the BF16 weight bit for bit; otherwise (or for a head the fused verify does not read) nothing is kept and + the head's own GEMM runs.""" + self._m3_target_head = None + if not _sharded_head(lm_head) or not m3_target_head_kernel.head_supported(lm_head): + return False + head = m3_target_head_kernel.prepare_head(lm_head) + self._m3_target_head = head + if not head.coded: + logger.warning( + f"[M3 target head] the coded shard {tuple(lm_head.weight.shape)} does not decode to the BF16 weight " + "bit for bit; the fused verify keeps the head's BF16 GEMM" + ) + return head.coded + + def _m3_coded_head(self, lm_head: nn.Module, rows: int): + """The coded target head for a fused verify of ``rows`` rows, or None for the head's own GEMM: above + ``TARGET_HEAD_MAX_ROWS`` rows (uncovered), on a shard whose coding failed its round trip (uncovered: the GEMM + is the exact path), and when the copy was not built for this weight (a miss).""" + path = TARGET_HEAD.path + if not TARGET_HEAD.covers(rows): + uncovered(path, f"{rows} verify rows (the head's own GEMM runs)") + return None + head = self._m3_target_head + if head is None or head.weight is not lm_head.weight: + missed( + path, + rows, + "the coded target head is not built for the head's current weight (prebuild)", + ) + return None + if not head.coded: + uncovered( + path, "the coded shard failed its bitwise round trip (the head's own GEMM runs)" + ) + return None + return head + + def _m3_force_vec( + self, + batch_size: int, + base_total: int, + max_total: int, + device, + kv_lens: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + ``FORCE_AL_PER_REQUEST`` (with ``kv_lens``): each request draws from the fixed-seed pool at its own key, its + KV length at this verify (its prompt plus what it has accepted so far), so its draws do not depend on the + server's history or on the other requests; every TP rank holds the same lengths, so the ranks agree. Without + lengths: exactly as _apply_force_accepted_tokens draws for a decode-only batch (one shared counter bump per + step, the (counter, slot) hash), the generic path's sequence. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: + return None + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + if FORCE_AL_PER_REQUEST and kv_lens is not None: + keys = kv_lens[:batch_size] + if keys.dtype != torch.int32: + keys = keys.to(torch.int32) + return m3_force_vec_req( + self._force_accept_rng_pool, + keys, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (batch_size - 1).bit_length()), + ) + return m3_force_vec( + self._force_accept_rng_pool, + self._force_accept_rng_counter, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (batch_size - 1).bit_length()), + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[M3Eagle3Draft]) -> bool: + """Build the streamed copy of the drafter's BF16 head shard that the fused draft loop reads (rebuilt from the + current weight values on every call), on every rank. False when the drafter is not one the loop serves.""" + self._m3_draft_head = None + if draft_model is None or not self._m3_drafter_supported(draft_model): + return False + self._m3_draft_head = m3_draft_step_kernel.prepare_head( + draft_model.lm_head, self.mapping.tp_rank, DRAFT_HEAD_CODED + ) + return True + + def _m3_drafter_supported(self, draft_model: M3Eagle3Draft) -> bool: + """The drafter the fused draft loop runs (``M3Eagle3Draft`` validates its shape at construction) over this + worker's tensor-parallel group with a bias-free vocab-sharded BF16 head without padding.""" + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if self._d2t is not None: + return False + lm_head = draft_model.lm_head + return _sharded_head(lm_head) and m3_draft_step_kernel.head_supported(lm_head) + + def _m3_draft_ar_norm( + self, + draft_model: M3Eagle3Draft, + x: torch.Tensor, + residual: torch.Tensor, + norm_key: str, + pushed: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """One drafter boundary of the fused draft loop: the multicast all-reduce + residual + RMSNorm in its drafter + form (``DRAFT_MC_AR``; bitwise the stock boundary), else the built-in path (``M3Eagle3Draft.ar_norm``). The + choice depends on shapes and a collectively created workspace only, so the TP ranks agree. ``pushed``: the + projection already pushed its rows into the all-reduce (``DRAFT_AR_PUSH``).""" + rows = x.shape[0] + path = DRAFT_MC_AR.path + if DRAFT_MC_AR.covers(rows): + workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) + out = None + if workspace is not None: + out = m3_mc_allreduce.m3_mc_allreduce_norm( + x.contiguous(), + residual.contiguous(), + draft_model.w[norm_key], + draft_model.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + drafter_norm=True, + carveout=DRAFT_MC_AR_CARVEOUT, + prepushed=2 if pushed else 0, + drafter_cluster=DRAFT_MC_AR_CLUSTER, + ) + if out is not None: + taken( + path, + rows, + "drafter boundary (cluster kernel)" + if DRAFT_MC_AR_CLUSTER + else "drafter boundary (row kernel)", + ) + if pushed: + taken(DRAFT_AR_PUSH_PATH, rows, "the projection pushed its rows (prepushed 2)") + elif DRAFT_AR_PUSH: + uncovered( + DRAFT_AR_PUSH_PATH, + "the boundary's projection did not push (the AR pushes x)", + ) + return out + if pushed: + raise RuntimeError( + "the drafter's rows were pushed into a boundary all-reduce that did not run" + ) + missed(path, rows, "the multicast all-reduce workspace or operand layout") + return draft_model.ar_norm(x, residual, norm_key) + + def _m3_arm_draft_push(self, weight: torch.Tensor, rows: int) -> None: + """``DRAFT_AR_PUSH``: let the projection ``weight``'s next low-M split-K call push its rows into the boundary + all-reduce that ``_m3_draft_ar_norm`` then runs (the multicast path only; shapes and a collectively created + workspace decide it, so the TP ranks agree).""" + if not (DRAFT_AR_PUSH and DRAFT_MC_AR.covers(rows)): + return + workspace = m3_mc_allreduce.get_workspace(self.mapping, weight.shape[0]) + if workspace is not None: + _draft_lowm_kernel.arm_push(weight, workspace, rows) + + def _m3_fused_loop_ok( + self, attn_metadata, spec_metadata, batch_size: int, draft_model: Optional[M3Eagle3Draft] + ) -> bool: + """Whether this step drafts with the fused loop: a decode-only step the fused verify or the seed prepared + (``_m3_first_gather_ids``) with the fused first draft inputs, a supported drafter, a (draft length, requests) + ``DRAFT_STEP`` covers, the prebuilt head and the exchange workspace. Host state only, so the TP ranks + agree.""" + if ( + self._m3_first_gather_ids is None + or not self._m3_fused_first_inputs + or draft_model is None + ): + return False + path = DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + uncovered( + path, "the drafter or its attention metadata is not one the fused loop serves" + ) + return False + draft_len = spec_metadata.runtime_draft_len + if not DRAFT_STEP.covers((draft_len, batch_size)): + uncovered(path, f"draft length {draft_len} at {batch_size} requests") + return False + head = self._m3_draft_head + if head is None or head.weight is not draft_model.lm_head.weight: + missed( + path, + batch_size, + "the drafter head's streamed copy is not built for its current weight (prebuild)", + ) + return False + if m3_draft_step_kernel.workspace(self.mapping, head.wf.device) is None: + capture_missed(path, None, "the draft-step exchange workspace") + return False + return True + + def _forward_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ): + """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop.""" + if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): + return self._m3_fused_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + return super()._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + + def _m3_fused_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model: M3Eagle3Draft, + draft_kv_cache_manager, + batch_size: int, + num_accepted_tokens, + ): + """Mirrors Eagle3OneModelWorker._forward_linear_draft_loop for a decode-only Eagle3 step, with each pass's + glue fused (see ``_m3_fused_loop_ok``): + + * the draft input -- the embedding lookup, the layer's input norm and hidden norm and their concat -- in one + launch (``m3_spec_tail.draft_input``; the second pass gathers the last accepted rows there); + * the drafter layer's attention and MLP, each followed by the all-reduce with the residual add and the next + RMSNorm (the post-attention norm, then the drafter's final norm); + * one kernel for the BF16 head over the head rows (pass 0: the last accepted row of each request), the + cross-rank argmax and the bookkeeping the in-tree loop does between passes: the draft token, the next pass's + input id and position, the draft KV lengths (and pass 0's sequence lengths) (``m3_draft_step``). + + The host metadata transitions are the in-tree loop's. The drafts land in columns 1.. of the step's next-input + tokens (column 0 is the last accepted token), which ``_forward_impl`` returns as they lie. + """ + draft_len = spec_metadata.runtime_draft_len + table = draft_model.embed_table + dev = table.device + head = self._m3_draft_head + gather_ids = self._m3_first_gather_ids + ids = inputs["input_ids"] + hidden = inputs["hidden_states"] + pos = inputs["position_ids"] + if ids.dtype != torch.int32: + ids = ids.to(torch.int32) + if pos.dtype != torch.int32: + pos = pos.to(torch.int32) + num_acc = num_accepted_tokens[:batch_size] + if num_acc.dtype != torch.int32: + num_acc = num_acc.to(torch.int32) + next_new = self._m3_next_new + if next_new is None: + # Seeded from the in-tree acceptance: the last accepted token is the first-pass input at the gather row. + next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=dev) + next_new[:, 0] = ids[gather_ids] + draft_tokens = next_new[:, 1:] + w = draft_model.w + eps = draft_model.rms_eps + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) as md: + for i in range(draft_len): + x, residual = m3_spec_tail.draft_input( + ids, + table, + hidden, + gather_ids if i == 1 else None, + w["in_norm"], + eps, + w["hidden_norm"], + eps, + ) + taken(DRAFT_INPUT.path, x.shape[0]) + rows = x.shape[0] + self._m3_arm_draft_push(w["o"], rows) + h = draft_model.attention(x, md, reduce=False) + h, residual = self._m3_draft_ar_norm( + draft_model, + h, + residual, + "post_norm", + pushed=_draft_lowm_kernel.take_pushed(h), + ) + self._m3_arm_draft_push(w["dn"], rows) + h = draft_model.mlp(h, reduce=False, silu_cpp=True) + h, to_save = self._m3_draft_ar_norm( + draft_model, + h, + residual, + "final_norm", + pushed=_draft_lowm_kernel.take_pushed(h), + ) + if draft_model.return_hidden_post_norm: + to_save = h + ids_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + pos_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + if not m3_draft_step_kernel.draft_step( + head, + h, + batch_size, + gather_ids if i == 0 else None, + self.mapping, + draft_tokens, + i, + ids_next, + pos, + pos_next, + md.kv_lens_cuda, + num_acc if i == 0 else None, + md._seq_lens_cuda if i == 0 else None, + draft_len, + ): + raise RuntimeError( + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" + ) + taken(DRAFT_STEP.path, batch_size, f"draft length {draft_len}") + ids, pos, hidden = ids_next, pos_next, to_save + if i == 0: + # The draft step set the device sequence lengths to 1 and rolled the KV lengths back. + md._seq_lens[:batch_size].fill_(1) + md.on_update() + has_kv_cache = md.kv_cache_manager is not None + if has_kv_cache and hasattr(md, "host_request_types"): + md.host_request_types[: md.num_contexts].fill_(1) + if has_kv_cache: + md.num_contexts = 0 + self._prepare_flash_mla_generation_layout(md, 0, batch_size) + md.update_for_spec_dec() + md.use_spec_decoding = False + else: + md.update_for_spec_dec() + self._m3_next_new_tokens = next_new + return draft_tokens + + +def _build_speculation(shell, model_config: ModelConfig): + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise ValueError( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + _set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return draft_config, draft, worker, use_separate + + +# -------------------------------------------------------------------------------------------------------------------- +# Runtime hooks and the shell +# -------------------------------------------------------------------------------------------------------------------- + + +def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: + """The step-loop hooks this target opts into (each behind its constant above). + + * ``pack_step_inputs``: a decode step's small input copies go out as one upload and one scatter at the graph + replay instead of a chain of copies the next graph waits on (~30 us of idle GPU at one request). + * ``skip_unchanged_h2d``: steady decode re-uploads slot indices, query lengths and page tables that rarely change. + * ``loop_on_execution_stream``: the overlap loop's own stream work and the forward on one stream, so step N+1's + inputs wait on graph N in program order rather than through two cross-stream event hops. + * ``merge_tuned_tactics_across_ranks``: tensor-parallel peers capture the same GEMM tactics (rank-local profiling + noise otherwise splits them, and the slowest rank sets every step). + * ``skip_greedy_rng_upload``: all-greedy speculative batches never read the Philox seeds. + * ``spec_step_fold``: the spec sampler's slot stores and the engine's next-step input gathers in one launch each. + * ``pack_parallel_scatter`` / ``pack_in_graph`` / ``pack_early_copy``: the pack's copies scatter in parallel, issued + by the decode graph's own prologue from a window uploaded ahead of the replay on a side copy stream. + * ``pack_spec_rng_uploads`` / ``spec_small_h2d``: the spec step's seeds, accepted counts and slot ids ride the pack + or stay on the device. + * ``sampler_d2h_side_stream`` / ``iteration_timing_one_event``: the sampler's store copies and the iteration timing + leave the main stream. + * MSA: identity selections for short sequences (and short/long graph families), the one-launch KV-length update, + unchanged block-table copies skipped, pool views built once, block-id tables copied from the page-index views. + """ + return RuntimeHooks( + first_content_emit_iters=FIRST_CONTENT_EMIT_ITERS, + mxfp8_pinned_tactics=True, + late_compile_detect=True, + pack_step_inputs=STEP_PACK_INPUTS, + pack_parallel_scatter=STEP_PACK_PARALLEL_SCATTER, + pack_in_graph=STEP_PACK_IN_GRAPH, + pack_early_copy=STEP_PACK_EARLY_COPY, + pack_step_epoch=STEP_PACK_EPOCH, + pack_prologue_op=m3_step_prologue if STEP_PACK_CPP_PROLOGUE else None, + skip_unchanged_h2d=STEP_SKIP_UNCHANGED_H2D, + h2d_changed_span=STEP_H2D_CHANGED_SPAN, + kv_lens_update_once=STEP_KV_LENS_UPDATE_ONCE, + spec_offsets_op=m3_spec_offsets if STEP_SPEC_OFFSETS_ONE_LAUNCH else None, + loop_on_execution_stream=STEP_LOOP_ON_EXECUTION_STREAM, + merge_tuned_tactics_across_ranks=STEP_MERGE_TUNED_TACTICS, + skip_greedy_rng_upload=STEP_SKIP_GREEDY_RNG_UPLOAD, + pack_spec_rng_uploads=STEP_PACK_SPEC_RNG_UPLOADS, + spec_small_h2d=STEP_SPEC_SMALL_H2D, + spec_step_fold=spec_fold, + sampler_d2h_side_stream=STEP_SAMPLER_D2H_SIDE_STREAM, + iteration_timing_one_event=STEP_ONE_TIMING_EVENT, + msa=MsaRuntimePolicy( + short_sequence_identity=MSA_SHORT_SEQUENCE_IDENTITY, + kv_update=_msa_kv_update_kernel.msa_kv_update if MSA_KV_UPDATE else None, + skip_unchanged_block_offsets=MSA_SKIP_UNCHANGED_BLOCK_OFFSETS, + reuse_pool_views=MSA_REUSE_POOL_VIEWS, + block_ids_from_views=MSA_BLOCK_IDS_FROM_VIEWS, + ), + ) + + +class ModelingV2MiniMaxM3Nvfp4Sm107Tp4( + DecoderModelForCausalLM[MiniMaxM3Tp4ModelingV2, PretrainedConfig] +): + """MiniMax-M3 NVFP4 / sm_107 / tp4, text only, with one-model Eagle3 or no speculation. + + The shell owns the vocab-parallel BF16 ``lm_head`` and the logits processor (``DecoderModelForCausalLM``); with + speculation, the drafter (``M3Eagle3Draft``, its weights from the speculative checkpoint through + ``load_draft_weights``) and the worker (``M3Eagle3Worker``) that verifies and drafts each step. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + assert torch.cuda.get_device_capability() == _SM, ( + f"this target is built for sm_{_SM[0]}{_SM[1]}, the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + # Under modeling_v2='require' a covered decode shape that cannot take its fast path raises. + set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + model_config = _text_model_config(model_config) + cfg = model_config.pretrained_config + spec_config = getattr(model_config, "spec_config", None) + if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: + # The one-model Eagle3 layer has no NVFP4 trtllm-gen decode kernel: keep its modules and the shared + # draft cache on FP8 while the target stays NVFP4. + model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 + super().__init__( + MiniMaxM3Tp4ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self.spec_config = spec_config + self.draft_model: Optional[M3Eagle3Draft] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[M3Eagle3Worker] = None + self.use_separate_draft_kv_cache = False + self.layer_idx = FINAL_HIDDEN_LAYER_ID + spec_fold = None + if spec_config is not None: + ( + self.draft_config, + self.draft_model, + self.spec_worker, + self.use_separate_draft_kv_cache, + ) = _build_speculation(self, model_config) + if STEP_SPEC_FOLD: + spec_fold = M3SpecStepFold( + in_forward=STEP_SPEC_FOLD_IN_FORWARD, + store_copy_wait=STEP_SPEC_STORE_COPY_SIDE, + cpp_tail=STEP_FOLD_TAIL_CPP, + ) + self._spec_fold = spec_fold + self.runtime_hooks = _runtime_hooks(spec_fold) + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights: Dict, weight_mapper=None) -> None: + """The drafter's checkpoint (``speculative_model``), loaded by the engine after the target's.""" + _weights.DRAFT_WEIGHTS.load(self.draft_model, weights) + + def set_guided_decoder(self, guided_decoder) -> bool: + """Mirrors ``SpecDecOneEngineForCausalLM.set_guided_decoder``: guided decoding with speculation is installed + on the worker.""" + if hasattr(self.spec_worker, "set_guided_decoder"): + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + def cache_derived_state(self) -> None: + """Build the targets and every decode kernel once the weights are real, before the executor's warmup; with + speculation also the drafter layer's full-context decode attention (``DRAFTER_FULLCTX``), the fused verify's + coded target head, the fused draft loop's head copy and the drafter's low-M GEMMs (``DRAFT_LOWM``).""" + super().cache_derived_state() + self.model.post_load_weights() + draft = self.draft_model + if draft is not None and DRAFTER_FULLCTX.keys: + path = DRAFTER_FULLCTX.path + installed = _drafter_attn_kernel.install( + draft, + {k: FULLCTX_DECODE_RECIPES[k] for k in DRAFTER_FULLCTX.keys}, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), + get_env_enable_pdl(), + fused_rope_mode=( + DRAFTER_ROPE_MODE + if DRAFTER_ROPE_FUSED + and hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_drafter_rope_fp8_kv") + else None + ), + on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), + on_taken=lambda rows: taken(path, rows), + on_uncovered=lambda reason: uncovered(path, reason), + on_missed=lambda rows, reason: missed(path, rows, reason), + ) + logger.info(f"[M3 prebuild] {path} installed on {installed} drafter attention layer(s)") + if self.spec_worker is not None: + self.spec_worker.prebuild(self.lm_head, draft) + if draft is not None: + # The drafter linears' low-M kernels, before the warmup. + start = time.perf_counter() + built = draft.prebuild_linears() + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] drafter low-M GEMMs built before warmup: %d calls, %.1f s" + % (built, time.perf_counter() - start) + ) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + """One engine step: the target forward, then the logits (or, with speculation, the worker's verify and + draft). Text only.""" + multimodal_params = kwargs.pop("multimodal_params", None) or [] + if any(p is not None and getattr(p, "multimodal_data", None) for p in multimodal_params): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + worker = self.spec_worker + capture = _plain_capture + if worker is not None: + capture = worker.capture_hidden_states + if spec_metadata is not None: + worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + capture=capture, + capture_target=worker.capture_target if worker is not None else None, + **kwargs, + ) + if getattr(attn_metadata, "csa2_remote_tail_mode", None) == "source": + # A context worker stopped at the encoder/decoder boundary: a disposable token for the context-only + # completion, and no draft worker. + return self.logits_processor.forward(hidden_states, self.lm_head, attn_metadata, False) + if spec_metadata is not None and spec_metadata.is_layer_capture(FINAL_HIDDEN_LAYER_ID): + capture(spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + if worker is None: + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], self.lm_head, attn_metadata, True + ) + spec_input_ids = input_ids if input_ids is not None else kwargs.get("orig_input_ids") + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if spec_input_ids is not None: + spec_input_ids = spec_input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens) + outputs = worker( + input_ids=spec_input_ids, + position_ids=spec_position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + spec_fold = self._spec_fold + if ( + spec_fold is not None + and spec_fold.in_forward + and attn_metadata.num_contexts == 0 + and isinstance(outputs, dict) + ): + spec_fold.store_in_forward(outputs) + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/__init__.py new file mode 100644 index 000000000000..b8c7426e3933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Activation ops of this target: the gated MLPs' fused SwiGLU.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..5bcdee3fbb82 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SiluAndMul(OpWrapper): + """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with + ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain + SwiGLU ``silu(gate) * up``. + + The MXFP8 dense MLPs (layers 0-2, intermediate 3072 per rank) and the MoE's shared expert (intermediate 768 per + rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 4608 per + rank) runs it plain. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=768, oai=True)), + Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=768, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=3072, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=3072, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=4608, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=4608, oai=False), + ), + ) + + note = """ + The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two + differ only by the fast sigmoid's last bits (the test's band). + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + gate, up = x.float().chunk(2, dim=-1) + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + a = 1.0 if alpha is None else alpha + b = 0.0 if beta is None else beta + return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" + + +m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/__init__.py new file mode 100644 index 000000000000..a7f7fee7a234 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention ops of this target: MSA block selection, the sparse decode and the KV-length update.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..ba008f493b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per +request count (``modeling.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 16 +KV_HEADS = 1 +HEAD_DIM = 128 +#: The split this target hands trtllm-gen at one request (modeling.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model +#: maximum the in-tree step passes otherwise. +CAPPED = 16384 +MODEL_MAX = 1 << 20 + + +def dense_attention_reference( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first + ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's + tokens in order). Returns fp32 [total_q, heads, head_dim].""" + batch, kv_heads, max_len, _ = keys.shape + heads = q.shape[1] + group = heads // kv_heads + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + pos = torch.arange(max_len, device=q.device) + with true_fp32_matmul(): + for b in range(batch): + full = int(seq_lens[b]) + rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] + limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) + valid = pos[None, :] < limit[:, None] # [qlen, max_len] + for h in range(kv_heads): + qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] + logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + probs = torch.softmax(logits, dim=-1) + out[ + b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group + ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) + return out + + +def gather_paged_kv( + kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool + (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + page = pool.shape[2] + pages = (int(seq_lens.max()) + page - 1) // page + slots = block_table[:, :pages].long() + k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] + v = pool[slots * subpages + 1] + batch, _, kv_heads, _, d = k.shape + + def order(t): + return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) + + return order(k), order(v) + + +class M3DenseDecode(OpWrapper): + """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len + ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. + + Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into + ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: + every CTA still covers its share of each request's actual length, which may exceed it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="two requests keep the model maximum (no cap at two requests)", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), + ), + ) + + note = """ + The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped + one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV + band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + minimax_m3_trtllm_gen_dense_decode( + q, + kv_cache_manager, + layer_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + max_seq_len=max_seq_len, + max_num_requests=max_num_requests, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). + + ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + if keys.dtype == torch.float8_e4m3fn: + q = fp8_e4m3_saturate(q) + return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> None: + # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. + assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" + # Each verify row attends a causal prefix of at least one key. + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + + +m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..d9e8ca93fff7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_qk_norm_rope_fp8.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head RMSNorm + partial NeoX RoPE of the packed [Q | K | V] into FP8, the bits of the BF16 +``fused_qk_norm_rope`` output cast to E4M3; optionally also storing the live tokens' K / V into the paged FP8 cache.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3DenseQkNormRopeFp8(OpWrapper): + """``qkv`` [T, (q + 2 kv) heads x 128] bf16 -> [T, same] E4M3: Q and K heads normed and roped in fp32, rounded to + BF16 and then to E4M3 (saturating), V cast; ``qkv`` untouched + (``trtllm::minimax_m3_sm107tp4_dense_qk_norm_rope_fp8``). + + With ``kv_buffers`` (the layer's HND FP8 pool [pages, 2, kv heads, 128, 128]), ``out_cache_loc`` (the step's int32 + cache slots) and ``num_live_tokens``, the kernel also stores the live tokens' K and V rows into the pool (the bytes + the MSA backend's KV write stores) and every row's Q into ``q_compact`` [T, q heads, 128] + (``..._dense_qk_norm_rope_fp8_kv``). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp4 heads (16 q, 1 kv)", spec=dict(m=1, q=16, kv=1)), + Cell( + why="a 4-request verify (16 rows) storing K / V", + spec=dict(m=16, q=16, kv=1, kv_write=True), + ), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=16, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=16, kv=1)), + ) + + note = """ + Bit-exact: `reference` runs the catalog's in-place BF16 ``fused_qk_norm_rope`` on a copy and casts every element to + E4M3, which is this kernel's rounding sequence (fp32 -> BF16 -> E4M3) for Q and K and the cast for V. + """ + + def __call__( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + use_gemma: bool, + kv_buffers: Optional[torch.Tensor] = None, + out_cache_loc: Optional[torch.Tensor] = None, + num_live_tokens: int = 0, + q_compact: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + args = ( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + use_gemma, + ) + if kv_buffers is not None: + return torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv( + *args, kv_buffers, out_cache_loc, num_live_tokens, q_compact + ) + return torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8(*args) + + def reference( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + use_gemma: bool, + kv_buffers: Optional[torch.Tensor] = None, + out_cache_loc: Optional[torch.Tensor] = None, + num_live_tokens: int = 0, + q_compact: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The FP8 copy (the K / V store into ``kv_buffers`` is checked by the test against the reference rows).""" + head_dim = qkv.shape[-1] // (num_heads_q + num_heads_k + num_heads_v) + out = qkv.clone() + fused_qk_norm_rope( + out, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + use_gemma, + False, + 0, + 0, + ) + return out.to(torch.float8_e4m3fn) + + +m3_dense_qk_norm_rope_fp8 = M3DenseQkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..ed4e3048d780 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention on the TRTLLM backend: RoPE, KV-cache append and causal attention.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class M3DraftAttention(OpWrapper): + """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the + backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's + NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` + ``num_tokens``); returns [num_tokens, q heads x 128] bf16. + + ``backend`` (``TrtllmAttention``) and ``metadata`` are the runtime's objects for the drafter layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="first draft pass of a 4-request verify (4 x 4 rows)", spec=dict(reqs=4, qlen=4)), + Cell(why="a later draft pass, one row per request", spec=dict(reqs=8, qlen=1)), + Cell(why="a context chunk's draft rows", spec=dict(reqs=1, context=512)), + ) + + note = """ + `reference` is exact (float64) causal softmax attention: Q rotated by the backend's NeoX RoPE at each token's + position, over the drafter layer's K / V rows read back from its FP8 cache (dequantized by the backend's + ``kv_scale_quant_orig``). What remains between the two is the FMHA's own arithmetic (Q's rounding to the + kernel's input precision, its P rounding and accumulation order), which the test bounds. + """ + + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors one KV head per rank and un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + assert backend.num_kv_heads == 1, "mirrors the one-KV-head-per-rank layout only" + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + keys = pool[pages, 0, 0].to(torch.float64).reshape(-1, dim)[:length] * kv_scale + values = pool[pages, 1, 0].to(torch.float64).reshape(-1, dim)[:length] * kv_scale + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] + rows = q[first : first + n] + x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,kd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) + + +m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..48543be75675 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q +out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3Fp8IndexerQkNormRope(OpWrapper): + """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and + writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, + then RoPE on its leading ``rotary_dim`` channels (NeoX). + + The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; + the fused layers' producer writes the index cache itself. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, one index head per rank", spec=dict(m=1, heads=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, heads=1)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=1)), + ) + + note = """ + `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 + before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of + values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. + """ + + def __call__( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + index_k_cache, + out_cache_loc, + num_heads_q, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + position_ids, + ) + + def reference( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, + then the E4M3 cast.""" + out = fused_qk_norm_rope( + qk.clone(), + num_heads_q, + 1, + 0, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) + return ( + q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), + k.to(torch.float8_e4m3fn), + ) + + +m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py new file mode 100644 index 000000000000..ad3abcadaae9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' full-context decode attention at TP4 on Rubin (``kernels/m3_fullctx_decode``): PrimTS split-KV +decode at the recipe ``modeling.FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) + +from ...kernels import m3_fullctx_decode as kernel +from .m3_dense_decode import dense_attention_reference, gather_paged_kv + +__all__ = ["m3_fullctx_decode"] + +HEADS = kernel.HEADS +KV_HEADS = kernel.KV_HEADS +HEAD_DIM = kernel.HEAD_DIM +#: The plan capacity (the per-request K/V length bound PrimTS sizes its policy for): dispatch's +#: FULLCTX_DECODE_MAX_KV_LEN, which the forward's prebuilt plans use. +MAX_KV_LEN = 1 << 21 + + +class M3FullCtxDecode(OpWrapper): + """Decode attention of every query row over its request's full causal context in the paged FP8 pool, written + into ``output`` (BF16 [T, heads, head_dim]) in place. + + ``recipe`` is ``modeling.FULLCTX_DECODE_RECIPES[(qlen, requests)]``: the PrimTS Q tile and the split-KV fanout. + ``own_reducer`` merges the splits on the BF16 op (``modeling.FULLCTX_DECODE_OWN_REDUCER``) rather than PrimTS's + FP16 reducer and a cast. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over an AgentX-length (165k) context", + spec=dict(kv_lens=(165525,), qlen=4), + ), + Cell( + why="one request, one draft-pass row, over a 165k context", + spec=dict(kv_lens=(165522,), qlen=1), + ), + Cell( + why="two requests verifying 1 + 3 rows at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165525, 120011), qlen=4), + ), + Cell( + why="two requests, one row each, at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165522, 120008), qlen=1), + ), + Cell( + why="one request verifying 1 + 3 rows over a short (4.1k) context: most splits are empty", + spec=dict(kv_lens=(4104,), qlen=4), + ), + ) + + note = """ + Same operands and arithmetic as the stock trtllm-gen call (E4M3 q, E4M3 K/V, P quantized to E4M3 per split, + fp32 accumulation): the split differs, so the result differs from stock as a stock split change does (the P + quantization is per split). Deterministic, and the same bits eager and in a CUDA graph. Held to the fp32 + reference within the FP8-KV band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + shape = (decode_query_len, int(seq_lens.shape[0])) + kernel.prebuild(q.device, {shape: recipe}, max_kv_len, own_reducer, get_env_enable_pdl()) + plan = kernel.plan_for(q.device, *shape, own_reducer) + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + k_pages, v_pages = kernel.kv_page_views(pool, subpages) + plan.run(fp8_e4m3_saturate(q), k_pages, v_pages, block_table, seq_lens, sm_scale, output) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V, q rounded to E4M3 first as the kernel runs it.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + return dense_attention_reference( + fp8_e4m3_saturate(q), keys, values, seq_lens, decode_query_len, sm_scale + ) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> None: + pool, _ = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + assert pool.dtype == torch.float8_e4m3fn, "an FP8 E4M3 pool" + assert tuple(pool.shape[1:]) == (KV_HEADS, kernel.PAGE, HEAD_DIM), ( + "the TP4 dense-layer pool geometry" + ) + assert q.shape[1:] == (HEADS, HEAD_DIM), "16 query heads of 128 per rank" + assert int(q.shape[0]) == int(seq_lens.shape[0]) * decode_query_len, ( + "q covers exactly the batch" + ) + if torch.cuda.is_current_stream_capturing(): + return # the length checks read the device: eager calls only + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + assert bool((seq_lens <= max_kv_len).all()), "a request longer than the plan capacity" + + +m3_fullctx_decode = M3FullCtxDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..f44d566310e7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``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_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over 20000 tokens: the early-issue kernel", + spec=dict(seq_lens=(20000,), qlen=4, early=True), + ), + Cell( + why="one request, one row, 65536 tokens: early issue, 512 blocks on 512 CTAs", + spec=dict(seq_lens=(65536,), qlen=1, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most CTAs lie past the length and are drained", + spec=dict(seq_lens=(129,), qlen=4, early=True), + ), + Cell( + why="one request longer than 512 blocks (the CTAs loop): early issue", + spec=dict(seq_lens=(70000,), qlen=4, early=True), + ), + Cell( + why="two requests: the prefetching kernel", + spec=dict(seq_lens=(20000, 3000), qlen=4, early=False), + ), + Cell( + why="eight requests of mixed lengths: the prefetching kernel", + spec=dict( + seq_lens=(300, 1500, 2049, 4096, 5000, 7777, 8192, 33001), qlen=4, early=False + ), + ), + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py new file mode 100644 index 000000000000..45103a9e289c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``kernels/m3_index_score_t2sc`` (T2SC: the sm_107 +TP2 / TP8 scorer as this target's copy, at its 1 index head per rank; ``modeling.INDEX_SCORE_T2SC``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_index_score_t2sc"] + +PAGE = 128 +HEAD_DIM = 128 + + +def compiled_scorer(key: tuple): + """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss + (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer_t2sc``).""" + from ...kernels.m3_index_score_t2sc import IndexScoreRunner + + if key not in IndexScoreRunner.kernel_cache: + IndexScoreRunner._compile(*key) + return IndexScoreRunner.kernel_cache[key] + + +class M3IndexScoreT2sc(OpWrapper): + """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the + positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of + ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is + written). FP8 or BF16 q and index-K, fp32 products and sums. + + ``early`` (the target's kernel at every request count it covers) issues the Q and first K tiles right after the + grid wait from a pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise + every older block of the CTA is prefetched. Neither, nor ``split_k``, changes a score. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request, 1 + 3 rows over 165521 tokens (AgentX ISL), 1 index head: 1536 CTAs, early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most of the 1536 CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request of 400000 tokens: 3125 blocks, so the 1536 CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request decoding one row over 65536 tokens", + spec=dict(seq_lens=(65536,), qlen=1, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request, 1 + 3 rows over 165521 tokens, 1 index head (TP4): 1272 CTAs (one wave), early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="one request of 129 tokens at 1272 CTAs: the early tiles of most CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="one request of 400000 tokens at 1272 CTAs: 3125 blocks, so the CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="two requests verifying 4 rows each over 165521 and 20000 tokens: 768 CTAs each, early issue", + spec=dict(seq_lens=(165521, 20000), qlen=4, heads=1, split_k=768, early=True), + ), + Cell( + why="two requests, one of 300 tokens and one of 262144 (2048 blocks: the 768 CTAs loop)", + spec=dict(seq_lens=(300, 262144), qlen=4, heads=1, split_k=768, early=True), + ), + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score_t2sc = M3IndexScoreT2sc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..0ce3c864d30f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MiniMax-M3 MSA attention backend's step: KV write, sparse top-k block selection, FMHA.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + +from .m3_sparse_decode_lc import dequant_nvfp4_pages + + +class M3MsaAttention(OpWrapper): + """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` + [T, heads * 128] (bf16, the step's live rows). + + * ``k`` / ``v`` given (the dense layers): written into the layer's paged main cache first (NVFP4 with + ``kv_scale_orig_quant``, else FP8 / BF16 as the pool is). None: a producer already wrote the caches. + * ``sparse``: the indexer scores the layer's index-K cache (already written) with ``idx_q`` and publishes the + top-k blocks; the FMHA attends those blocks. Otherwise the FMHA attends the whole page table. + + The FMHA dispatches on the step's metadata (context rows: the MSA prefill FMHA; generation rows: the MSA decode + FMHA, the dense layers' trtllm-gen decode at their capped KV split). ``backend`` and ``metadata`` are the runtime's + objects for this layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dense layer, a prefill chunk, FP8 pool", + spec=dict(layer="dense", phase="context", tokens=1024), + ), + Cell( + why="dense layer, a decode step of 4 requests", + spec=dict(layer="dense", phase="generation", reqs=4), + ), + Cell( + why="sparse layer, a prefill chunk beyond 16 blocks, NVFP4 pool", + spec=dict(layer="sparse", phase="context", tokens=4096), + ), + Cell( + why="sparse layer, a mixed step (context + generation rows)", + spec=dict(layer="sparse", phase="mixed", tokens=512, reqs=2), + ), + ) + + note = """ + `reference` is exact (float64) softmax attention over the cache rows the step attends, read back from the paged + cache the call wrote (or a producer wrote before it) and dequantized as the FMHA dequantizes them: FP8 rows as + E4M3, NVFP4 rows as E2M1 x E4M3 block scale x the per-tensor ``kv_scale_quant_orig``. A sparse layer's top-k blocks + are re-selected from ``idx_q`` and the index-K cache with the backend's rule (per-block max over the causal + keys, the query's own block forced in, ties to the lower block). On an FP8 pool the FMHA takes Q as a plain E4M3 + cast, which the reference applies too. What remains between the two is the FMHA's own arithmetic (its P + rounding and accumulation order), which the test bounds. + """ + + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors one KV head per rank, + the target's layout. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + assert backend.num_kv_heads == 1, "mirrors the one-KV-head-per-rank layout only" + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + if nvfp4: + keys = dequant_nvfp4_pages( + pool[:, 0].view(torch.uint8), scale_pool[:, 0], pages, quad_scales=False + ) + values = dequant_nvfp4_pages( + pool[:, 1].view(torch.uint8), scale_pool[:, 1], pages, quad_scales=True + ) + keys, values = keys * k_global, values * v_global + else: + keys = pool[pages, 0, 0].to(torch.float64) + values = pool[pages, 1, 0].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + qpos = cached[row] + torch.arange(n, device=q.device) + allowed = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + if sparse: + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = _selected_blocks( + idx64[first : first + n], idx_keys, allowed, qpos, page, config + ) + allowed &= chosen.repeat_interleave(page, dim=1)[:, :length] + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = torch.einsum("thd,kd->thk", q64[first + lo : first + hi], keys) * sm_scale + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + (ids >= n_valid - config.local_blocks) & (ids < n_valid) + ) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen + + +m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py new file mode 100644 index 000000000000..5a3d1b64ac58 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer, with its constant loads staged +before the PDL grid-dependency wait and an early dependents trigger (``kernels/minimaxM3/producerPw.cu``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +__all__ = ["m3_producer"] + +HEAD_DIM = 128 +ROTARY_DIM = 64 +PAGE = 128 +#: Query heads per rank at TP4 (MiniMax-M3: 64 query heads, 1 KV head and 1 index head per rank). +HEADS_Q = 16 +#: NVFP4 block: 16 values share one E4M3 scale. +FP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _norm_rope( + x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(FP4_BLOCK) + + +def scale_index(within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return (within_page // 4) * 32 + cols * 4 + within_page % 4 + + +class M3Producer(OpWrapper): + """The NVFP4 QKV + indexer producer of a sparse layer: ``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert`` for one + KV head, one index head, head dim 128 and rotary dim 64, with the same arguments and outputs bit for bit. + + ``packed`` [T, (heads_q + 4) * 128] BF16 holds each token's q heads, k, v, index-q and index-k. Every head but v + gets the Gemma RMSNorm with its weight and NeoX RoPE on its first 64 dims (``rotary_cos_sin[position]``). Returns + (q FP8 [T, heads_q, 128], index-q FP8 [T, 1, 128]); writes the token's slot ``out_cache_loc[t]`` (skipped when + negative or past the pool): k and v as NVFP4 (16-value blocks, E4M3 scales against ``kv_quant_scale[1]`` / + ``[2]``, the V scales in the 4-token swizzle) and index-k as FP8. ``early`` triggers the dependent kernel on + entry (its outputs are read only after the dependent's own grid wait). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token (c1 without speculation)", spec=dict(tokens=1, scale=1.0)), + Cell(why="a c1 verify: 4 tokens", spec=dict(tokens=4, scale=3.0)), + Cell(why="a c4 verify: 16 tokens, small activations", spec=dict(tokens=16, scale=1e-3)), + Cell( + why="a c8 verify: 32 tokens, large activations, a skipped and an out-of-pool slot", + spec=dict(tokens=32, scale=40.0, bad_slots=True), + ), + Cell(why="the widest decode step: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="a prefill chunk of 1000 tokens, early trigger off", + spec=dict(tokens=1000, scale=1.0, early=False), + ), + ) + + note = """ + The kernel is the native producer's arithmetic, compiled with the library's --use_fast_math, so it is held to that + kernel bit for bit on every output and cache byte. The torch reference states the math in fp32: the FP8 outputs are + held to it within one E4M3 step, and the NVFP4 K/V, decoded, within one E2M1 step of their block scale. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if prewait_step_inputs: + # This target's copy (producerPw4): the same outputs bit for bit, with the position ids, slots and RoPE + # rows read before its grid wait (none of them may be written by a PDL predecessor). + return tuple( + torch.ops.trtllm.minimax_m3_sm107_tp4_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + True, + ) + ) + q, index_q = torch.ops.trtllm.minimax_m3_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + ) + return q, index_q + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", + "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change + them.""" + t = packed.shape[0] + x = packed.float().view(t, heads_q + 4, HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = rope[:, 0], rope[:, 1] + q = torch.stack( + [_norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + ) + return { + "q": q, + "k": _norm_rope(x[:, heads_q], k_weight, cos, sin, eps), + "v": x[:, heads_q + 1], + "index_q": _norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), + "index_k": _norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), + } + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> None: + t = packed.shape[0] + # The RoPE row is read at the position with no bound. + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + # Two tokens writing one slot race. + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" + + +m3_producer = M3Producer() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..f92a3e6efaa1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3QkNormRopeFp8(OpWrapper): + """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded + once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place + entry. + + The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps + of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp4 heads (16 q, 1 kv)", spec=dict(m=1, q=16, kv=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=16, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=16, kv=1)), + ) + + note = """ + `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from + fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. + """ + + def __call__( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + + def reference( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" + out = fused_qk_norm_rope( + qkv.clone(), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return out.to(torch.float8_e4m3fn) + + +m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..ef79579c0c21 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into +the paged caches.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3QkvInsert(OpWrapper): + """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the + partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and + index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. + + ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool + [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 + [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp4 + heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), + Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), + ) + + note = """ + The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` + returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + common = ( + num_heads, + num_kv_heads, + num_index_heads, + head_dim, + rotary_dim, + eps, + q_norm, + k_norm, + index_q_norm, + index_k_norm, + rotary_cos_sin, + position_ids, + ) + if kv_scale is not None: + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common + ) + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, index_k_cache, out_cache_loc, *common + ) + + def reference( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of + Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as + packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" + m = packed.shape[0] + sizes = [ + num_heads * head_dim, + num_kv_heads * head_dim, + num_kv_heads * head_dim, + num_index_heads * head_dim, + head_dim, + ] + q, k, v, iq, ik = packed.float().split(sizes, dim=-1) + cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) + half = rotary_dim // 2 + + def norm_rope(x, w, heads): + x = x.reshape(m, heads, head_dim) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) + x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] + c, s = cos.unsqueeze(1), sin.unsqueeze(1) + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) + + return { + "q": norm_rope(q, q_norm, num_heads), + "index_q": norm_rope(iq, index_q_norm, num_index_heads), + "k": norm_rope(k, k_norm, num_kv_heads), + "v": v.reshape(m, num_kv_heads, head_dim), + "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], + } + + +m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..76841cca732d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SelectBlocks(OpWrapper): + """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), + with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the + row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` + backs the result head-major so ``permute(1, 0, 2)`` is contiguous. + + The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="eight decode rows of 17..128 blocks", + spec=dict(total_q=8, n_blocks=128, mixed=True), + ), + Cell( + why="one decode row at the 128-block edge", + spec=dict(total_q=1, n_blocks=128, mixed=False), + ), + Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), + ) + + note = """ + Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. + """ + + def __call__( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + return torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output + ) + + def reference( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + """Plain torch selection; ``head_major_output`` changes the backing only.""" + heads, n_blocks, total_q = max_score_kv.shape + out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) + blocks = torch.arange(n_blocks, device=max_score_kv.device) + for t in range(total_q): + nv = int(n_valid_blocks[t]) + for h in range(heads): + s = max_score_kv[h, :, t].clone().double() + s[blocks >= nv] = float("-inf") + forced = (blocks < min(init_blocks, nv)) | ( + (blocks >= nv - local_blocks) & (blocks < nv) + ) + s[forced] = float("inf") + order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] + ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) + out[t, h, : ids.numel()] = ids + return out + + +m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..3d076c8181d6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). +MAX_BLOCKS = 9216 +#: Blocks the selection keeps per (query, KV head). +TOPK = 16 +#: Scores the forcing substitutes: initial blocks rank above local ones, and both above any finite score. +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 + +# M3's checkpoint: no forced initial blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=1, init_blocks=0, local_blocks=1) + + +class M3SelectLong(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query's top-16 blocks. + + ``scores`` is the per-KV-head block maxima [num_kv_heads, num_blocks, + total_q] the CuTe DSL index scorer wrote (amax over the KV head's index + heads). Per (query, KV head) row the op keeps the ``n_valid_blocks[query]`` + leading blocks (clamped to ``[0, num_blocks]``), forces the last + ``local_blocks`` of them (counted from the *unclamped* valid count) to + ``LOCAL_SCORE`` and the first ``init_blocks`` to ``INIT_SCORE`` -- local wins + where both apply -- and selects the 16 largest, ties to the smaller block. + A pick whose forced score is -inf is emitted as -1. The ids are written + ascending with a -1 tail. + + Same contract as ``torch.ops.trtllm.minimax_m3_select_blocks`` (the in-tree + ``select_blocks_from_maxscore``) and bit-identical to it, except on a row + where more than 1024 valid blocks tie at the 16th-largest score: this op + breaks that tie to the smallest blocks as the contract says, the in-tree op + does not. The test checks both against ``reference``. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="32K-context engine (257 worst-case k-tiles), Eagle3 verify at one request: 4 query rows", + spec=dict(num_blocks=257, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="decode without speculation at batch 8 on a 2048-block engine: the 2-key-per-thread path at its edge", + spec=dict(num_blocks=2048, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="1M-context engine at batch 8 x qlen 4: rows on all three key widths, short rows and padding rows", + spec=dict(num_blocks=8193, total_q=32, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 1024 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + note = """ + Certified at one KV head per rank (TP4) with M3's forcing (`init_blocks=0`, + `local_blocks=1`). The kernel takes any head count and forcing; neither is + driven here. + + `scores` may be strided (the op reads it by stride); the target hands over + the scorer's contiguous view, which is what the cells drive. Rows of at + most 128 blocks are within the kernel's range, but the target keeps those on + the in-tree op, so they are not certified here. + + Ordering is by the descending-score key `orderedKey`: a positive NaN ranks + above +inf, +0 above -0, and a negative NaN below -inf -- which, unlike + -inf, is emitted as a block id rather than -1. The cells drive positive + NaN only. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sm107_tp4_select_long( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + fast_tail, + ) + return out + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and + ``trigger_phase`` and ``fast_tail`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() # [total_q, num_blocks] + forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where( + in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max + ) + k = min(TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) + ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> None: + """The inputs this op takes and answers wrongly. + + Negative forcing counts are accepted and push the forced ranges past the + row instead of failing; the in-tree op this one stands in for refuses + them, so a caller relying on this one would silently diverge from it. + An ``out`` overlapping ``scores`` is overwritten while it is still read. + """ + assert init_blocks >= 0 and local_blocks >= 0, ( + f"init_blocks ({init_blocks}) and local_blocks ({local_blocks}) must be non-negative; the " + "op accepts negative counts and selects as if the forced ranges were elsewhere" + ) + storage = scores.untyped_storage() + assert not (storage.data_ptr() <= out.data_ptr() < storage.data_ptr() + storage.nbytes()), ( + "out must not alias scores: the op writes it while scores are still read" + ) + + +m3_select_long = M3SelectLong() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py new file mode 100644 index 000000000000..9f79f4b69517 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_select_long import _M3, MAX_BLOCKS, M3SelectLong + + +class M3SelectLongDsl(OpWrapper): + """``m3_select_long``'s contract, cells and reference, run by ``kernels/m3_select_long_dsl``: the same selection + bit for bit as the 1024-thread ``minimax_m3_sm107_tp4_select_long`` (also on rows where more than 2048 blocks tie + at the 16th-largest score), from one 384-thread CTA per row. ``fast_tail`` (the 1024-thread kernel's warp-0 rank) + does not apply and is ignored. Returns ``out``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="32K-context engine (257 worst-case k-tiles), Eagle3 verify at one request: 4 query rows", + spec=dict(num_blocks=257, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="decode without speculation at batch 8 on a 2048-block engine: the 2-key-per-thread path at its edge", + spec=dict(num_blocks=2048, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="1M-context engine at batch 8 x qlen 4: rows on all three key widths, short rows and padding rows", + spec=dict(num_blocks=8193, total_q=32, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 1024 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> torch.Tensor: + """``m3_select_long``'s reference: the selection is the same bit for bit.""" + return M3SelectLong.reference( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + fast_tail, + ) + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> None: + """``m3_select_long``'s domain.""" + M3SelectLong.is_valid( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + fast_tail, + ) + + note = """ + Certified at one KV head per rank (TP4, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`), and bit for bit against the + 1024-thread `m3_select_long` on every cell. The kernel takes any head count + and forcing; other forcing is not driven here. + + The kernel is compiled per (PDL launch, early trigger, trigger phase) with + symbolic shapes and strides; the target compiles its keys before warmup + (`kernels/m3_select_long_dsl.prebuild`), and a call outside them compiles + on first use. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> torch.Tensor: + from ...kernels.m3_select_long_dsl import select_long_dsl + + if scores.shape[1] > MAX_BLOCKS: + raise RuntimeError( + f"m3_select_long_dsl supports at most {MAX_BLOCKS} blocks, got {scores.shape[1]}" + ) + select_long_dsl( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + +m3_select_long_dsl = M3SelectLongDsl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..04d76f241b29 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This target's block-sparse decode attention (Triton, with the one-launch long-context route).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="identity selection, decode without speculation at batch 8: the Triton producer + merge", + spec=dict( + seq_lens=(1, 128, 129, 700, 2048, 1500, 64, 0), decode_query_len=1, identity=True + ), + ), + Cell( + why="identity selection, Eagle3 verify at one request (T=4): one block per producer program", + spec=dict(seq_lens=(1900,), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 4 (T=16): the last row count below the MQ producer", + spec=dict(seq_lens=(300, 2048, 4, 1000), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 8 (T=32): the multi-query producer, padding rows", + spec=dict( + seq_lens=(2048, 130, 5, 0, 1777, 900, 2000, 0), decode_query_len=4, identity=True + ), + ), + Cell( + why="selected blocks, Eagle3 verify at batch 2: routed to the one-launch long-context kernel", + spec=dict(seq_lens=(9000, 40000), decode_query_len=4, identity=False), + ), + Cell( + why="selected blocks, decode without speculation at batch 8: the long-context route at T=8", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + identity=False, + ), + ), + ) + + note = """ + Certified only as the target calls it: one KV head of 16 query heads, + head_dim 128, FP8 q, packed NVFP4 cache, BF16 output, `num_topk_chunks` + left to the dispatcher. The function also takes BF16 / FP8 caches, other + head geometries and an explicit split-K factor (the in-tree tests drive + those); none is certified here. + + The MXFP8 hand-off each route publishes is certified on the long-context + kernel's own entry (`m3_sparse_decode_lc`); the Triton merge writes it with + the same recipe (`_store_mxfp8_head`). + """ + + def __call__( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + kv_scale=kv_scale, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + num_topk_chunks=num_topk_chunks, + identity_blocks=identity_blocks, + ) + return output + + def reference( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + """``output`` as the call is supposed to write it: float64 block-sparse attention over the + dequantized NVFP4 blocks (``sparse_attention_reference``), rounded once to ``output``'s dtype. + + Only the packed NVFP4 cache with one KV head is mirrored -- the target's + only configuration. ``kv_scale``, ``num_topk_chunks`` and + ``identity_blocks`` do not change what is computed. + """ + assert k_block_scale is not None and k_paged.shape[1] == 1, ( + "mirrors the one-KV-head NVFP4 cache only" + ) + attn = sparse_attention_reference( + q, + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.to(output.dtype) + + def is_valid( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) + + +m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py new file mode 100644 index 000000000000..76fc68fa5903 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py @@ -0,0 +1,436 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch long-context sparse decode over the selected NVFP4 blocks (16-CTA clusters). + +Also the plain-torch block-sparse attention and MXFP8 hand-off mirrors the +Triton sparse decode entry (``m3_sparse_decode``) reuses: both compute the same +thing over the same cache. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The TP4 geometry the kernel is compiled for. +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +#: NVFP4 elements per E4M3 block scale. +SF_VEC = 16 +#: MXFP8 elements per UE8M0 scale, and the o_proj width the hand-off is laid out for. +MX_BLOCK = 32 +WIDTH = NUM_HEADS * HEAD_DIM + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, block_scale: torch.Tensor, pages: torch.Tensor, quad_scales: bool +) -> torch.Tensor: + """Dequantize KV head 0 of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, 0].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, 0].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """Block-sparse decode attention over one KV head of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages(k, k_block_scale, pages, quad_scales=False).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True).reshape(-1, HEAD_DIM) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, + sf [num_rows_alloc * WIDTH // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) + num_sf_k_tiles = WIDTH // 128 + sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +class M3SparseDecodeLc(OpWrapper): + """Sparse decode attention, split-K merge and MXFP8 hand-off of one TP4 decode step in one launch. + + One cluster of 16 CTAs per query token: CTA ``r`` attends the token's + ``r``-th selected block (``topk[tok, r]``), and the cluster merges the 16 + partials over distributed shared memory into the BF16 row ``out[tok]``. With + ``q8`` / ``sf`` it also writes the row's MXFP8 copy and swizzled UE8M0 + scales for the o_proj GEMM. All operands are the byte views + ``m3_attn_lc.m3_lc_attn`` hands over (E4M3 ``q``, packed E2M1 pages, E4M3 + block scales). Returns ``out``. + + Per block: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16 + (exact), QK and PV on FP16 tensor cores with FP32 accumulation, the softmax + probabilities rounded to FP16 for PV, log2-domain softmax with ``ex2.approx``. + ``linear_softmax`` picks the denominator recurrence (the two agree to + rounding); the target picks it per batch shape. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation, one request: T=1, log-domain recurrence, MXFP8 hand-off", + spec=dict(seq_lens=(9000,), decode_query_len=1, linear_softmax=False, mxfp8=True), + ), + Cell( + why="Eagle3 verify at one request: T=4, the log-domain recurrence the target keeps for local batch 1", + spec=dict(seq_lens=(20000,), decode_query_len=4, linear_softmax=False, mxfp8=True), + ), + Cell( + why="Eagle3 verify at two requests: T=8 on the linear denominator", + spec=dict(seq_lens=(5000, 12345), decode_query_len=4, linear_softmax=True, mxfp8=True), + ), + Cell( + why="Eagle3 verify at batch 8 (T=32): short requests with -1 tails and zero-length padding rows", + spec=dict( + seq_lens=(16384, 2049, 130, 0, 7777, 4, 30000, 0), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + ), + ), + Cell( + why="decode without speculation at batch 8: T=8 rows of distinct requests", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + linear_softmax=False, + mxfp8=True, + ), + ), + Cell( + why="no hand-off buffers (the variant the call site takes past 128 rows), driven at T=8", + spec=dict(seq_lens=(5000, 12345), decode_query_len=4, linear_softmax=True, mxfp8=False), + ), + ) + + note = """ + Certified: the four instantiated variants (`linear_softmax` x hand-off on / + off), all with `early_trigger=True` -- the op refuses `False`. The + hand-off cells are certified bit-exactly against `mxfp8_handoff_reference` + applied to the op's own BF16 rows; the BF16 rows themselves within the band + the test derives. + + `topk` is read for the first `min(16, ceil(kv_len / 128))` entries of each + token; a -1 among them skips that CTA's block, entries past them are never + read. `use_pdl` only changes the launch attribute. The hand-off assumes + `num_sf_k_tiles = 16` (a 2048-wide o_proj input) and buffers with at least + as many rows as the step; `MERGE_MXFP8` sizes them at 128 rows. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + bt_prewait: bool = False, + ) -> torch.Tensor: + # merge_once: this target's copy of the op (minimaxM3Sm107Tp4/sparseDecodeLc4.cu), whose split-K merge forms the + # 16 weights once per CTA (same values, same order: bitwise the image op's output). + lc = ( + torch.ops.trtllm.minimax_m3_sm107_tp4_sparse_decode_lc + if merge_once + else torch.ops.trtllm.minimax_m3_sparse_decode_lc + ) + # prefetch: up to two read-only tensors the TP4 op pulls into L2 (a hint; the image op takes none). + # next_kv: (this layer's id buffer, the next sparse layer's, its K / V pools and block scales), TP4 op only. + extra = () + if merge_once and (prefetch or next_kv): + extra = tuple(prefetch or ()) + (None,) * (2 - len(prefetch or ())) + if next_kv: + extra += tuple(next_kv) + lc( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + decode_query_len, + sm_scale, + linear_softmax, + use_pdl, + early_trigger, + *extra, + **({"identity": True} if identity else {}), + **({"identity_step": True} if identity_step else {}), + **({"one_div": True} if one_div else {}), + # bt_prewait: the TP4 op stages the block-table row in shared memory before its grid wait (same output). + **({"bt_prewait": True} if bt_prewait and merge_once else {}), + ) + return out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + bt_prewait: bool = False, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 2048] the op is supposed to write: ``sparse_attention_reference`` in float64, + rounded once. ``linear_softmax``, ``use_pdl``, ``early_trigger``, ``one_div`` and ``bt_prewait`` do not change + what is computed; the hand-off buffers are checked separately with ``mxfp8_handoff_reference``. + """ + attn = sparse_attention_reference( + q.view(-1, NUM_HEADS, HEAD_DIM), + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.reshape(out.shape).to(out.dtype) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + bt_prewait: bool = False, + ) -> None: + """The inputs this op takes and answers wrongly (``m3_attn_lc.lc_applies`` is the target's form of the + layout half). + + The kernel addresses a page's token rows at a fixed 64-byte (data) and + 8-byte (scale) pitch and loads them 16 / 8 bytes at a time, so only the + page stride is read from the views. It reads the selected block ids, + the block table and the hand-off buffers unchecked. Reads ``topk`` and + ``seq_lens`` back to the host: a test-only cost. + """ + for name, t in (("k", k), ("v", v)): + assert t.dim() == 4 and t.shape[1] == 1 and t.shape[2:] == (PAGE, HEAD_DIM // 2), ( + f"{name} must be the [pages, 1, 128, 64] byte view of one KV head" + ) + assert t.stride(3) == 1 and t.stride(2) == HEAD_DIM // 2, ( + f"{name} rows must be dense 64-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0, ( + f"{name} pages must be 16-byte aligned" + ) + for name, t in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + assert t.dim() == 4 and t.stride(3) == 1 and t.stride(2) == HEAD_DIM // SF_VEC, ( + f"{name} rows must be dense 8-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 8 == 0 and t.stride(0) % 8 == 0, ( + f"{name} pages must be 8-byte aligned" + ) + assert q.data_ptr() % 16 == 0 and q.stride(0) % 16 == 0 and q.stride(1) % 16 == 0, ( + "q rows are loaded 16 bytes at a time and must be 16-byte aligned" + ) + assert k_global_scale.numel() == 1 and v_global_scale.numel() == 1, ( + "global scales are single values" + ) + total_q = q.shape[0] + assert ( + total_q % decode_query_len == 0 and seq_lens.numel() >= total_q // decode_query_len + ), ( + "every query token must belong to a request: seq_lens is read at tok // decode_query_len unchecked" + ) + seq = seq_lens.tolist() + num_cols = block_table.shape[1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk[tok, : min(TOPK, nblocks)] + assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( + f"token {tok}: selected block ids must be -1 or below the request's block count; the block " + "table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + f"num_sf_k_tiles must be {WIDTH // 128}: the hand-off rows are {WIDTH} wide and a smaller " + "value overlaps them silently" + ) + assert ( + q8.numel() >= total_q * WIDTH + and sf is not None + and sf.numel() >= ((total_q + 127) // 128 * 128 * (WIDTH // MX_BLOCK)) + ), "the hand-off buffers must hold every row of the step: they are written unchecked" + + +m3_sparse_decode_lc = M3SparseDecodeLc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..43166b472eca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The MSA pool's page: one sparse block is one 128-token KV page. +TOKENS_PER_BLOCK = 128 + + +class MsaKvUpdate(OpWrapper): + """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. + + ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its + own buffers at the head of every captured speculative decode step. One + single-CTA launch writes, for ``batch`` requests and ``total_q`` query + tokens (all int32): + + * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, + ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); + * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` + attends up to, through ``block_table``; + * ``n_valid_blocks[t]``, when given, the block count top-k selection is + bounded by. + + It returns None; the cells read the four buffers. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation at one request: qo = 1, every buffer written", + spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="decode without speculation at the captured batch 8", + spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", + spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), + ), + Cell( + why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", + spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), + ), + Cell( + why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", + spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), + ), + Cell( + why="positions crossing a page boundary: the slot moves to the previous page of the request", + spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), + ), + ) + + note = """ + The kernel covers any batch up to the block size's loop (it strides over + 256 threads), and any page size; only `tokens_per_block = 128`, the MSA + pool's page, is certified. + + The op does not bounds-check the request index a token reads + (`q_batch_row`) or the block-table column it derives: both come from the + metadata's own staging, which keeps them in range by construction, and an + out-of-range value reads whatever memory sits there. `is_valid` states both. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.minimax_m3_msa_kv_update( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + + def reference( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], + n_valid_blocks[:total_q])`` as the op is supposed to write them. + + The in-tree metadata's own torch re-derivation + (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook + call), statement for statement, returning new tensors rather than + writing the buffers. The buffer arguments only say which outputs exist: + a ``None`` buffer gives a ``None`` result. + """ + page = int(tokens_per_block) + kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) + cu_kv = None + if cu_kv_lens is not None: + cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) + cu_kv[1:] = torch.cumsum(kv_true, 0) + qbr = q_batch_row[:total_q].to(torch.long) + qo_dev = seq_lens[:batch] + pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(block_table.shape[1]) - 1 + ) + slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + n_valid = None + if n_valid_blocks is not None: + n_valid = torch.div( + (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" + ).to(torch.int32) + return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid + + def is_valid( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + """The inputs this op takes and answers wrongly. + + Every operand but the block table is read as a dense run from its first + element, so a strided view reads the wrong elements. The request index of + each token and the table row it selects are not bounds-checked. Reads + ``q_batch_row`` back to the host: a test-only cost. + """ + for name, t in ( + ("kv_lens", kv_lens), + ("kv_lens_staged", kv_lens_staged), + ("seq_lens", seq_lens), + ("q_batch_row", q_batch_row), + ("q_intra", q_intra), + ("msa_seq_lens", msa_seq_lens), + ("out_cache_loc", out_cache_loc), + ("cu_kv_lens", cu_kv_lens), + ("n_valid_blocks", n_valid_blocks), + ): + assert t is None or t.is_contiguous(), ( + f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" + ) + rows = q_batch_row[:total_q] + assert total_q == 0 or ( + int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) + ), ( + "q_batch_row must index a request below batch (and a block-table row): the op reads the " + "request's lengths and table row unchecked" + ) + + +msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/__init__.py new file mode 100644 index 000000000000..5da86fe24bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective-communication entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..d67b040e8bfd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams + + +class M3AllReduce(OpWrapper): + """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the + plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: + ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; + returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the + communicator's per call. + + The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps + beyond the multicast kernel's rows. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="tp4, a prefill chunk", spec=dict(world=4, m=2048, hidden=6144)), + Cell(why="tp4, a 64-row generation step", spec=dict(world=4, m=64, hidden=6144)), + Cell( + why="tp4, plain sum of eight drafter rows", + spec=dict(world=4, m=8, hidden=6144, plain=True), + ), + ) + + note = """ + The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then + adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum + to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. + """ + + def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): + return comm(x, all_reduce_params=params) + + def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): + """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" + total = torch.stack([x.float() for x in xs]).sum(0) + if params is None: + return total.to(xs[0].dtype) + updated = (params.residual.float() + total).to(xs[0].dtype) + h = updated.float() + normed = ( + h + * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) + * params.norm_weight.float() + ) + return normed.to(xs[0].dtype), updated + + +m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..208cb25aa72e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,329 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode layer boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer + fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; + only the stored residual is rounded to BF16. + + The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are + summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic + and the normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + pre = total.to(torch.bfloat16).float() + residual.float() + updated = pre.to(torch.bfloat16) + u = pre.double() if drafter else updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +class M3McAllreduceNorm(OpWrapper): + """One TP group's boundary: ``updated = bf16(residual + sum_r bf16(x_r + x2_r))``, ``normed = RMSNorm(updated)``. + + Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls + its own view ``ws_uc`` (re-arming the words it read) and rotates the per-CTA call parities in ``flags``; the + workspace is the ``McAllReduceWorkspace`` bucket of ``maxt`` tokens. All ranks must make the same calls in the + same stream order. ``prepushed`` 2 means a producer already pushed ``x`` (and launched this grid after its + pushes); 3 that two producers pushed ``x`` into exchange set A and ``x2`` into set B. Returns + ``[normed, updated]`` or, with ``quant``, also the MXFP8 copy of ``normed`` (E4M3 rows, swizzled UE8M0). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The tp4 decode boundaries (hidden 6144, 4 ranks, eps 1e-6). Buckets 8 / 16 / 48 are separate workspaces; the + # 8-token bucket runs epilogue option 7 at 1, 2, 7, 8 tokens and 15 at 3..6, the others 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token at the attention boundary, x only, bucket 8 option 7", + spec=dict(m=1, maxt=8, opts=7, x2=False, quant=False, prepushed=0), + ), + Cell( + why="MoE boundary at 4 tokens: shared + routed summed in-kernel, MXFP8 copy for the next QKV, option 15", + spec=dict(m=4, maxt=8, opts=15, x2=True, quant=True, prepushed=0), + ), + Cell( + why="a full 8-token bucket with the second addend", + spec=dict(m=8, maxt=8, opts=7, x2=True, quant=False, prepushed=0), + ), + Cell( + why="attention boundary whose rows the o_proj GEMV already pushed (prepushed 2)", + spec=dict(m=8, maxt=8, opts=7, x2=False, quant=False, prepushed=2), + ), + Cell( + why="routed engine pushed set A, shared down projection set B (prepushed 3), with the MXFP8 copy", + spec=dict(m=6, maxt=8, opts=15, x2=True, quant=True, prepushed=3), + ), + Cell( + why="12 tokens pushed by the o_proj GEMV into the 16-token bucket", + spec=dict(m=12, maxt=16, opts=15, x2=False, quant=False, prepushed=2), + ), + Cell( + why="the 9..13-token MoE kernel's sum into the 16-token bucket, MXFP8 copy", + spec=dict(m=16, maxt=16, opts=15, x2=False, quant=True, prepushed=0), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=30, maxt=48, opts=15, x2=False, quant=False, prepushed=0), + ), + Cell( + why="the 48-token bucket full, second addend and MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, x2=True, quant=True, prepushed=0), + ), + # Every other row count of the decode boundaries (modeling.MC_ALLREDUCE): the bucket and epilogue options the + # target picks for it (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): + # modeling.ATTN_PUSH / modeling.MOE_PUSH. + *( + Cell( + why=f"{m} tokens prepushed: a pushed row count between the ones above", + spec=dict( + m=m, + maxt=8 if m <= 8 else 16, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1] if m <= 8 else 15, + x2=False, + quant=m <= 8, + prepushed=2, + ), + ) + for m in range(1, 17) + if m not in (8, 12) + ), + # The Eagle3 drafter boundaries (modeling.DRAFT_MC_AR): option 16 (the row kernel's drafter form) on 23 / 31 + # at every draft-pass row count 1..8 (pass 0: 4 x requests, later passes: requests), x only. + *( + Cell( + why=f"Eagle3 drafter boundary at {m} rows (opts {23 if m in (1, 2, 7, 8) else 31})", + spec=dict( + m=m, + maxt=8, + opts=23 if m in (1, 2, 7, 8) else 31, + x2=False, + quant=False, + prepushed=0, + drafter=True, + ), + ) + for m in range(1, 9) + ), + # The same boundaries on the 8-CTA cluster kernel (opts bit 128, modeling.DRAFT_MC_AR_CLUSTER, launched without + # the carveout: 23 | 128), alone and fed by the drafter projection's push (prepushed 2, DRAFT_AR_PUSH). + *( + Cell( + why=f"Eagle3 drafter boundary on the cluster kernel at {m} rows, prepushed {pre}", + spec=dict( + m=m, + maxt=8, + opts=23 | 128, + x2=False, + quant=False, + prepushed=pre, + drafter=True, + ), + ) + for pre in (0, 2) + for m in range(1, 9) + ), + ) + + note = """ + The epilogue option bits (1 early residual / weight load, 2 st.async cluster exchange, 4 early parity rotation, + 8 max carveout) are bit-identical by construction; only the four (bucket, option) pairs the target selects are + built. `prepushed` 1 (pushed, grid wait kept) is accepted by the op but no caller passes it. `pdl` changes only + launch overlap and is driven off in the cells. + + `normed` is gated to one BF16 rounding step: the kernel's RMS statistic is a warp / cluster tree of FP32 sums + and its `rsqrtf` is approximate, so a row whose scaled value lands next to a BF16 midpoint can round either + way. `updated` and the MXFP8 copy of the op's own `normed` are bit-exact. + """ + + def __call__( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_sm107_tp4_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + None, + None, + maxt, + quant, + prepushed, + opts, + cap_out, + cap_w, + cap_eps, + ) + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant, drafter) + + def is_valid( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> None: + h = residual.shape[-1] + # The op checks norm_w's element count, not its layout: a strided view is read as packed memory. + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + # The workspace is addressed as [set][parity][maxt][world][H / 2] words and never bounds-checked: a smaller + # buffer (another bucket's, or one built for fewer ranks) is polled and re-armed out of range. + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + # Flags are per (token, cluster CTA) of the bucket; the op checks the count, not that they belong to the + # workspace. Flags of another bucket hold another call history, and the parity they select is not the one the + # peers push into: the poll then waits on the other half forever or reads last call's rows. + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + + +m3_mc_allreduce_norm = M3McAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..0d457022c38f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize + multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode MoE boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_mc_allreduce_norm import boundary + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class M3McFinalizeAllreduceNorm(OpWrapper): + """The MoE boundary with the routed experts' finalize folded in. + + Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, + ``scales`` fp32 or bf16), then the group's all-reduce, residual add and RMSNorm exactly as + ``m3_mc_allreduce_norm``: same workspace, flags and call-order contract. Returns ``[normed, updated]`` or, with + ``quant``, also the MXFP8 copy of ``normed``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The tp4 target takes this path from 9 tokens (top-4 of 128 experts, fp32 routing weights, the shared expert + # always present): the 16- and 48-token buckets, epilogue option 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="the smallest finalize call (9 tokens, 16-token bucket) with the next QKV's MXFP8 copy", + spec=dict(m=9, maxt=16, opts=15, topk=4, quant=True, dropped=0), + ), + Cell( + why="a full 16-token bucket with dropped expert slots", + spec=dict(m=16, maxt=16, opts=15, topk=4, quant=False, dropped=5), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=24, maxt=48, opts=15, topk=4, quant=False, dropped=0), + ), + Cell( + why="the 48-token bucket full, MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, topk=4, quant=True, dropped=3), + ), + # Every other row count of the unfinalized MoE (modeling.MC_FINALIZE): the bucket and options the target picks. + *( + Cell( + why=f"{m} tokens: a finalize row count between the ones above", + spec=dict( + m=m, + maxt=16 if m <= 16 else 48, + opts=15, + topk=4, + quant=m % 2 == 0, + dropped=m % 4, + ), + ) + for m in range(9, 49) + if m not in (9, 16, 24, 48) + ), + ) + + note = """ + BF16 routing weights and the 8-token bucket are accepted by the op (option 7 there) but the target only enters + this path at 9 tokens and up, with the built-in MoE's fp32 weights; neither is a cell. The rest is as in + `m3_mc_allreduce_norm`'s note. + """ + + def __call__( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_sm107_tp4_mc_allreduce_norm( + fc2, + shared, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + 0, + opts, + ) + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + ) -> None: + h = residual.shape[-1] + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + # The kernel gathers fc2[perm] without a bound: an index past the FC2 rows reads whatever follows them. + lo, hi = int(perm.min()), int(perm.max()) + assert lo >= -1 and hi < fc2.shape[0], ( + f"perm spans [{lo}, {hi}]; entries must be -1 (dropped) or a row of fc2 (< {fc2.shape[0]})" + ) + + +m3_mc_finalize_allreduce_norm = M3McFinalizeAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/__init__.py new file mode 100644 index 000000000000..e9f4d8e20b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedding ops of this target: the generic token-embedding gather.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..3060c17acdff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="a 48-row decode verify, int64 ids", + spec=dict(rows=48, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ) + + note = """ + An exact gather: `reference` must match bit for bit. + """ + + def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) + + def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] + return rows * valid.unsqueeze(-1).to(table.dtype) + + +m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/__init__.py new file mode 100644 index 000000000000..9cc374bdf1cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMM entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..4fb5af737dc0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3Bf16Linear(OpWrapper): + """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. + + The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), + o_proj (this rank's input columns), gate_up and down projections run it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=2304, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=2048)), + Cell(why="drafter gate_up, one row", spec=dict(m=1, n=9216, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=4608)), + ) + + note = """ + cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x, w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return (x.float() @ w.float().t()).to(x.dtype) + + +m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py new file mode 100644 index 000000000000..1fb7a1d759ae --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense MLP's swigluoai activation as one C++ launch with PDL and an early release (DENSE_SWIGLU_CPP).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# The dense layers' swigluoai parameters, which the kernel hard-codes (modeling_minimaxm3._build_swiglu_oai_dense_mlp). +LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 + + +class M3DenseSwiglu(OpWrapper): + """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, + rounded once to bf16: trtllm::minimax_m3_sm107_tp4_swiglu_oai (minimaxM3Sm107Tp4/swigluOai.cu), bit for bit the + Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the + grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights + during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even + half width; ``__call__`` returns None otherwise or for other activation parameters (the caller runs the Triton + op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP4 dense MLP shard: gate | up 6144 wide (d = 3072), at the decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP4 dense MLP at 1 row", spec=dict(rows=1, d=3072)), + Cell(why="the TP4 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=3072)), + Cell(why="the TP4 dense MLP at 8 rows", spec=dict(rows=8, d=3072)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the + alpha-scaled ex2.approx sigmoid, div.full), checked over every bf16 gate and every bf16 up pattern. + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float] = LIMIT, + alpha: Optional[float] = ALPHA, + beta: Optional[float] = BETA, + use_pdl: bool = True, + ) -> Optional[torch.Tensor]: + if ( + (limit, alpha, beta) != (LIMIT, ALPHA, BETA) + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 4 != 0 + or x.data_ptr() % 4 != 0 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp4_swiglu_oai(x, out, bool(use_pdl)) + return out + + def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_dense_swiglu = M3DenseSwiglu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py new file mode 100644 index 000000000000..6fb8b1572d22 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 x BF16 -> BF16 low-M GEMM of the Eagle3 drafter's linears at fixed FlashInfer tactics (SM10x).""" + +from __future__ import annotations + +from typing import Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: (N, K) of the TP4 drafter shards -> ((max rows, kind, tactic), ...), rows ascending. Kind "d" is FlashInfer's direct +#: SIMT kernel with tactic (block, outputs per block, rows per block); kind "s" is its split-K tensor-core kernel with +#: tactic (mma_m, mma_n, split_k, ab_stages). A row count is served by the first entry whose max rows holds it; the +#: rows modeling.DRAFT_LOWM covers are only the ones measured (the fastest tactic at no larger error vs fp64 than +#: cuBLAS). Rubin SM107 TP4 sweep (mclk locked, >= 3x L2 of rotating weights, CUDA-graph timing), us cuBLAS -> +#: low-M at 1 / 4 rows: qkv 8.15 -> 5.90 / 8.32 -> 6.71, o 5.36 -> 3.40 / 5.44 -> 3.66, gate_up 12.93 -> 9.93 / +#: 13.03 -> 12.05, down 7.68 -> 6.16 / 7.64 -> 6.15. The fc is replicated (the TP2 shape): its TP2 tactic and row. +TACTICS: Mapping[Tuple[int, int], tuple] = { + # qkv + (2304, 12288): ((1, "d", (256, 2, 1)), (4, "s", (64, 8, 4, 10))), + # o_proj + (6144, 2048): ((1, "d", (128, 4, 1)), (4, "s", (64, 8, 1, 10))), + # gate_up + (9216, 6144): ((1, "d", (128, 2, 1)), (4, "s", (64, 8, 1, 6))), + # down + (6144, 4608): ((4, "s", (64, 8, 2, 10)),), + # the Eagle3 fc over the three captured hidden states + (6144, 18432): ((1, "d", (384, 2, 1)),), +} +#: Rows each (N, K) is certified at (a subset of what its table would serve). +ROWS: Mapping[Tuple[int, int], Tuple[int, ...]] = { + (2304, 12288): (1, 4), + (6144, 2048): (1, 4), + (9216, 6144): (1, 4), + (6144, 4608): (1, 4), + (6144, 18432): (1,), +} + + +def tactic_for(n: int, k: int, m: int) -> Optional[Tuple[str, tuple]]: + """(kind, tactic) of an [m, k] x [n, k]^T call, or None when the table does not serve it.""" + for rows, kind, tactic in TACTICS.get((n, k), ()): + if m <= rows: + return kind, tactic + return None + + +class M3DraftLowM(OpWrapper): + """``out`` [M, N] BF16 = ``x`` [M, K] BF16 @ ``weight`` [N, K]^T BF16, FP32 accumulation. + + ``kind`` / ``tactic`` pick FlashInfer's kernel: "d" runs ``run_direct_dense`` (SIMT, one dot product per output in + K order), "s" runs ``run_splitk_dense`` (tcgen05 MMA, K split over ``split_k`` CTAs and summed in FP32). Both + round once to BF16. The kernel is compiled on its first call per (M, N, K, tactic, PDL); the target compiles every + covered call before warmup (kernels/m3_draft_lowm.prebuild). + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter qkv, 1 rows: d tactic", spec=dict(nk=(2304, 12288), m=1)), + Cell(why="drafter qkv, 4 rows: s tactic", spec=dict(nk=(2304, 12288), m=4)), + Cell(why="drafter o_proj, 1 rows: d tactic", spec=dict(nk=(6144, 2048), m=1)), + Cell(why="drafter o_proj, 4 rows: s tactic", spec=dict(nk=(6144, 2048), m=4)), + Cell(why="drafter gate_up, 1 rows: d tactic", spec=dict(nk=(9216, 6144), m=1)), + Cell(why="drafter gate_up, 4 rows: s tactic", spec=dict(nk=(9216, 6144), m=4)), + Cell(why="drafter down, 1 rows: s tactic", spec=dict(nk=(6144, 4608), m=1)), + Cell(why="drafter down, 4 rows: s tactic", spec=dict(nk=(6144, 4608), m=4)), + Cell(why="drafter fc, 1 rows: d tactic", spec=dict(nk=(6144, 18432), m=1)), + ) + + note = """ + The table is the TP4 drafter's: its shapes are the per-rank shards of the Eagle3 layer at 4-way tensor + parallelism (32 query and 2 KV heads per rank, intermediate 9216 per rank) and the replicated fc, at the b1 + decode rows (1 row in draft passes 1-2, 4 in pass 0). Certified at those shapes and rows only; FlashInfer's + kernels take other shapes, which are not measured here. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + if kind == "d": + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + DirectTactic, + run_direct_dense, + ) + + run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) + elif own_splitk: + # This target's copy of the split-K kernel (kernels/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # it also pushes the rows into a drafter boundary all-reduce (modeling.DRAFT_AR_PUSH). + from ...kernels.m3_lowm_splitk import ( + SplitKTactic, + run_splitk_dense, + run_splitk_dense_push, + ) + + if ar is not None: + run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) + else: + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + else: + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + return out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """bf16(x @ weight^T) summed in fp64 and rounded once; ``out``, ``kind``, ``tactic`` and ``use_pdl`` only + choose where and how the op writes.""" + return (x.double() @ weight.double().t()).to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> None: + """The kernels index ``x`` and ``out`` densely: a strided view is read as if it were packed.""" + assert x.dim() == 2 and x.is_contiguous(), "x must be a contiguous [M, K] matrix" + assert weight.dim() == 2 and weight.is_contiguous(), ( + "weight must be a contiguous [N, K] matrix" + ) + assert out.is_contiguous() and tuple(out.shape) == (x.shape[0], weight.shape[0]), ( + "out must be a contiguous [M, N] matrix" + ) + assert ar is None or (own_splitk and kind == "s" and len(ar) == 5), ( + "ar (the boundary all-reduce push) runs only on this target's split-K copy" + ) + + +m3_draft_lowm = M3DraftLowM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..9702cc0778a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class M3Fp32Linear(OpWrapper): + """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the + fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies + (the engine leaves torch's default).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), + Cell( + why="a generation step above the tensor-core router's 48 rows", + spec=dict(m=64, e=128, k=6144), + ), + ) + + note = """ + Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against + the true fp32 reference. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x.to(torch.float32), w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return x.float() @ w.t() + + +m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..c6ccde3e5ae4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 x MXFP8 GEMV (tcgen05, block-scaled, CuTe DSL) on an activation its producer already quantized.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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 tp4, the GEMV's tuned shapes. +QKV_SPARSE = (2560, 6144) +QKV = (2304, 6144) +O_PROJ = (6144, 2048) +SHARED_GATE_UP = (1536, 6144) +SHARED_DOWN = (6144, 768) +DENSE_GATE_UP = (6144, 6144) +DENSE_DOWN = (6144, 3072) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "shared gate_up": SHARED_GATE_UP, + "shared down": SHARED_DOWN, + "dense gate_up": DENSE_GATE_UP, + "dense down": DENSE_DOWN, +} + + +def served_cells(named: tuple[Cell, ...], x_bf16: bool, max_m: int) -> tuple[Cell, ...]: + """A cell for every (projection, rows) the GEMV serves up to ``max_m`` rows that ``named`` does not list. + + Certification is per projection: a projection's split-K, ring and MMA width differ, so each (N, K, M) the + decode dispatch can send here is its own cell. + """ + listed = {(cell.spec["nk"], cell.spec["m"]) for cell in named} + return tuple( + Cell(why=f"{name}, {m} rows", spec=dict(nk=nk, m=m)) + for name, nk in PROJECTIONS.items() + for m in range(1, max_m + 1) + if (nk, m) not in listed and _gemv.serves(nk[0], nk[1], m, x_bf16) + ) + + +def gemv_reference( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """bf16(x @ dequant(weight)^T) for fp32 activation values x [M, K], summed in fp64 and rounded once.""" + return (x.double() @ mxfp8_dequant(weight, weight_scale).double().t()).to(torch.bfloat16) + + +def gemv_is_valid(weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + n, k = weight.shape + # The kernel reads both from data_ptr() as dense buffers at the layout it was compiled for. + assert weight.dtype == torch.float8_e4m3fn and weight.is_contiguous(), ( + "weight must be contiguous e4m3 [N, K]" + ) + assert weight_scale.numel() == n * k // 32, ( + "weight_scale must be the 128x4-swizzled UE8M0 buffer [N * K / 32]" + ) + + +class M3MxGemv(OpWrapper): + """bf16 [M, N] = xq [M, K] (e4m3, 128x4-swizzled UE8M0 scales ``xsf``) x weight [N, K]^T (e4m3, swizzled + ``weight_scale``), M <= 48; None when the shape is not served (the caller falls back). + + ``trigger_early``: release the PDL dependents before the epilogue. ``cfg``: a (split, ring[, contig, nacc, a_tmem, + st_async]) tuning override; default the table's. ``use_pdl``: default the environment's. ``ar`` = (multicast + workspace address, flags address, rank, world, bucket tokens, mode): also push the rows into the next boundary + all-reduce's exchange (then ``op.last_pushed`` is True). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The projections' (N, K) at the token counts each build differs in: the MMA width (8 / 16 / 32 / 48 columns) a + # call uses is the smallest that holds M, and the shared gate_up takes another split at 2 tokens and above 8. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 4)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell(why="sparse-layer QKV, 8 tokens", spec=dict(nk=QKV_SPARSE, m=8)), + Cell( + why="sparse-layer QKV, 48 verify rows: the widest MMA", spec=dict(nk=QKV_SPARSE, m=48) + ), + Cell(why="dense-attention QKV, one token", spec=dict(nk=QKV, m=1)), + Cell(why="dense-attention QKV, 13 rows: the 16-column MMA", spec=dict(nk=QKV, m=13)), + Cell(why="o_proj, one token (split 2)", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 24 rows: the 32-column MMA", spec=dict(nk=O_PROJ, m=24)), + Cell( + why="shared gate_up, one token (A from shared memory)", + spec=dict(nk=SHARED_GATE_UP, m=1), + ), + Cell(why="shared gate_up, 2 tokens: its split-4 config", spec=dict(nk=SHARED_GATE_UP, m=2)), + Cell( + why="shared gate_up, 12 rows: the wide split-4 config", + spec=dict(nk=SHARED_GATE_UP, m=12), + ), + Cell(why="shared down, one token (unsplit, K 768)", spec=dict(nk=SHARED_DOWN, m=1)), + Cell(why="shared down, 8 tokens", spec=dict(nk=SHARED_DOWN, m=8)), + Cell(why="dense gate_up, one token", spec=dict(nk=DENSE_GATE_UP, m=1)), + Cell(why="dense down, 8 tokens", spec=dict(nk=DENSE_DOWN, m=8)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling.MX_GEMV_ROWS). + CELLS = CELLS + served_cells(CELLS, x_bf16=False, max_m=48) + + note = """ + The kernel sums in fp32 in split-K order where `reference` sums in fp64: the bf16 outputs agree within the band the + test derives. Not covered: the all-reduce push (`ar`), which needs a multicast exchange; `cfg` overrides beyond the + table (the target passes none). Shapes outside the table, M > 48, and 17..32 rows at K < 1024 return None. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv(xq, xsf, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The dequantized product; the launch options are ignored.""" + return gemv_reference(mxfp8_dequant(xq, xsf), weight, weight_scale) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + m, k = xq.shape + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous(), ( + "xq must be contiguous e4m3 [M, K]" + ) + # A linear-layout [M, K / 32] scale buffer is read as the swizzled one: wrong scales, no error. + assert xsf.numel() >= (m + 127) // 128 * 128 * (k // 32), ( + "xsf must be the 128x4-swizzled scale buffer (rows padded to 128), not the linear one" + ) + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv = M3MxGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..f8b2ab00217f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 GEMV (CuTe DSL) on a bf16 activation, quantized to MXFP8 in the kernel's prologue.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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, + SHARED_DOWN, + SHARED_GATE_UP, + gemv_is_valid, + gemv_reference, + served_cells, +) + +__all__ = ["m3_mx_gemv_bf16"] + + +class M3MxGemvBf16(OpWrapper): + """bf16 [M, N] = mxfp8_quantize(x) x weight [N, K]^T for a contiguous, 16-byte aligned bf16 ``x`` [M, K]; the + quantization is mxfp8_quantize's (UE8M0 scale per 32 = the smallest power of two >= amax / 448). None when the + shape is not served. The other arguments as ``m3_mx_gemv``'s. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The prologue quantizes one 32-element block per thread and round: the projections at the token counts the + # target sends here (no producer-quantized activation), up to the widest each serves. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 16 rows: the most its split-4 prologue serves", + spec=dict(nk=QKV_SPARSE, m=16), + ), + Cell(why="dense-attention QKV, 8 tokens", spec=dict(nk=QKV, m=8)), + Cell(why="o_proj, one token", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 32 rows: the 32-column MMA", spec=dict(nk=O_PROJ, m=32)), + Cell(why="shared gate_up, one token", spec=dict(nk=SHARED_GATE_UP, m=1)), + Cell(why="shared gate_up, 2 tokens: its split-4 config", spec=dict(nk=SHARED_GATE_UP, m=2)), + Cell(why="shared down, 8 tokens (K 768)", spec=dict(nk=SHARED_DOWN, m=8)), + Cell( + why="dense gate_up, 8 tokens: the most split 2 at K 6144 serves", + spec=dict(nk=DENSE_GATE_UP, m=8), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling.MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=32) + + note = """ + As `m3_mx_gemv`. The in-kernel quantization is mxfp8_quantize's bit for bit except where the kernel's + rcp.approx(448) and an exact 1/448 put a block's amax on different sides of a power of two (not observed on the + cells). Not covered: `ar`, `cfg` overrides. Returns None for rows beyond 8 rounds of the prologue per CTA. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv_bf16(x, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The MXFP8 round trip of x times the dequantized weight; the launch options are ignored.""" + return gemv_reference(mxfp8_quant_dequant(x), weight, weight_scale) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv_bf16 = M3MxGemvBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..0b422a7167e5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned + +from .m3_mxfp8_linear import mxfp8_reference + + +class M3Mxfp8Gemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight + [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key + and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The decode dispatch hands it the boundary / merge MXFP8 activation on rows the GEMV does not serve: the + # unsplit shared down projection at 17..32 rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=768)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=768)), + ) + + note = """ + Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + """The dequantized product; ``tune`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..d7a70d725189 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned + +_BLOCK = 32 + + +def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" + pad_r = (rows + 127) // 128 * 128 + pad_c = (cols + 3) // 4 * 4 + s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) + return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] + + +def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, k = x.shape + s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) + factor = torch.pow(2.0, (s - 127).to(torch.float32)) + return ( + x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) + ).view(rows, k) + + +def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: + """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" + with true_fp32_matmul(): + return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=2560, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=2048, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=3072, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=1536, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=2048, path="tuned")), + ) + + note = """ + Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` + quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and + multiplies dequantized operands in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + if tuned_quantize_gemm: + return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + if flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, output_dtype + ) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are + ignored.""" + m, k = x.shape + blocks = x.float().view(m, k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(-1, keepdim=True) + exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) + scale = torch.pow(2.0, exp) + q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale + with true_fp32_matmul(): + out = q.view(m, k) @ _dequant(weight, weight_scale).t() + return (out * alpha.float()).to(output_dtype) + + +m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py new file mode 100644 index 000000000000..ee781b16ff42 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM at a fixed FlashInfer tactic (MX-TAC): the same choice on every rank and every start.""" + +from __future__ import annotations + +from typing import Dict, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + +from .m3_mxfp8_linear import mxfp8_reference + +# One runner per output dtype (a runner holds no per-call state). +_RUNNERS: Dict[Tuple[type, torch.dtype], object] = {} + + +def _runner(dtype: torch.dtype): + runner = _RUNNERS.get((FlashInferMXFP8GemmRunner, dtype)) + if runner is None: + runner = _RUNNERS[(FlashInferMXFP8GemmRunner, dtype)] = FlashInferMXFP8GemmRunner(dtype) + return runner + + +class M3Mxfp8PinnedGemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + swizzled UE8M0 scales) and an MXFP8 weight [N, K] + (swizzled scales) on FlashInfer's MXFP8 GEMM at ``tactic``, no AutoTuner; ``output_dtype`` out.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The decode graph's MXFP8 GEMMs on rows the GEMVs do not serve: the shared down projection on its boundary + # activation (17..32 rows) and the bf16 projections past the GEMV rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=768)), + Cell(why="packed QKV + index projection, 48 rows", spec=dict(m=48, n=4992, k=6144)), + ) + + note = """ + The MXFP8 product (fp32 accumulation, one rounding to the output dtype): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. The tactic only picks the kernel. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tactic: int, + ) -> torch.Tensor: + return _runner(output_dtype).forward([act, act_scale, weight, weight_scale], tactic=tactic) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tactic: int, + ) -> torch.Tensor: + """The dequantized product; ``tactic`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_pinned_gemm = M3Mxfp8PinnedGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py new file mode 100644 index 000000000000..3f35af65048a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 activation quantize at a fixed tactic (MX-TAC), ahead of ``m3_mxfp8_pinned_gemm``.""" + +from __future__ import annotations + +from typing import Dict, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8QuantizeRunner + +_RUNNERS: Dict[torch.dtype, object] = {} + + +def _runner(dtype: torch.dtype): + runner = _RUNNERS.get(dtype) + if runner is None: + runner = _RUNNERS[dtype] = MXFP8QuantizeRunner(dtype) + return runner + + +class M3Mxfp8PinnedQuantize(OpWrapper): + """A bf16 activation [M, K] to MXFP8: (e4m3 [M, K], swizzled UE8M0 scales), on the runner's ``tactic``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="packed QKV + index projection input, 48 rows", spec=dict(m=48, k=6144)), + Cell(why="shared expert down projection input, 17 rows", spec=dict(m=17, k=768)), + ) + + note = """ + Bit-exact: `reference` is the stock MXFP8 quantize (``torch.ops.trtllm.mxfp8_quantize``, swizzled scales) of the + same input; every tactic computes the same block maxima and roundings. + """ + + def __call__(self, x: torch.Tensor, tactic: int) -> Tuple[torch.Tensor, torch.Tensor]: + return _runner(x.dtype).forward([x], tactic=tactic) + + def reference(self, x: torch.Tensor, tactic: int) -> Tuple[torch.Tensor, torch.Tensor]: + """The stock quantize; ``tactic`` is ignored.""" + return torch.ops.trtllm.mxfp8_quantize(x, True) + + +m3_mxfp8_pinned_quantize = M3Mxfp8PinnedQuantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py new file mode 100644 index 000000000000..f8deb858c843 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, +and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). + +The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values +``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 +marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a +per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: +per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 +code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved +it decodes to the BF16 weight bit for bit. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_target_head"] + +#: Verify rows per call (requests x (draft length + 1)). +MAX_ROWS = 8 +#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). +CTA_WARPS = 16 +#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. +TARGET_SHARD = 50016 +HIDDEN = 6144 +#: Codes per window: 0..14 decode, 15 marks an exception. +WINDOW = 15 +#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. +MAX_BASE = 254 - WINDOW +#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). +ROWS_PER_CHUNK = 4096 + + +def _pad_rows(weight: torch.Tensor) -> torch.Tensor: + n = weight.shape[0] + tiles = (n + 15) // 16 + if tiles * 16 == n: + return weight + return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) + + +def _fragment_view(w: torch.Tensor) -> torch.Tensor: + """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" + n, k = w.shape + # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) + return ( + w.view(n // 16, 2, 8, k // 16, 2, 4, 2) + .permute(0, 3, 2, 5, 4, 1, 6) + .reshape(n // 16, k // 16, 32, 8) + ) + + +def _exponents(weight: torch.Tensor) -> torch.Tensor: + return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF + + +def _window_base(hist: torch.Tensor) -> int: + """The base of the 15-exponent window holding the most weights (the lowest such base).""" + window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) + return int(window.argmax()) + + +def encode(weight: torch.Tensor, base: Optional[int] = None): + """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's + own window. + + Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile + << 16 | k), exc_w fp32 [E, 2] ({w, 0})). + """ + w = _pad_rows(weight) + n, k = w.shape + tiles, ks = n // 16, k // 16 + assert ks % 4 == 0, "K must be a multiple of 64" + bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF + expo = (bits >> 7) & 0xFF + if base is None: + base = _window_base(torch.bincount(expo.flatten(), minlength=256)) + code = expo - base + exc = (code < 0) | (code >= WINDOW) + code = torch.where(exc, torch.full_like(code, WINDOW), code) + sm = ((bits >> 8) & 0x80) | (bits & 0x7F) + sm = torch.where(exc, torch.zeros_like(sm), sm) + # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. + rows, cols = exc.nonzero(as_tuple=True) + order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) + rows, cols = rows[order], cols[order] + exc_w = torch.stack( + [ + w[rows, cols].float(), + torch.zeros((rows.numel(),), device=w.device), + ], + dim=1, + ) + exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) + exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) + exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( + torch.int32 + ) + # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. + smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) + cf = ( + _fragment_view(code.to(torch.uint8).view(n, k)) + .view(tiles, ks // 4, 4, 32, 8) + .to(torch.int32) + ) + # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. + cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) + for p in range(4): + cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) + cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 + # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. + seg01 = ( + smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) + ) # [.., 2, 32, 16] + seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) + packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] + return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() + + +def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): + """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" + n, _ = weight.shape + assert rows_per_chunk % 16 == 0 + hist = torch.zeros(256, dtype=torch.int64, device=weight.device) + for r in range(0, n, rows_per_chunk): + hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) + base = _window_base(hist) + packed, offs, poss, ws = [], [], [], [] + n_exc = 0 + for r in range(0, n, rows_per_chunk): + p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) + packed.append(p) + offs.append(off[:-1] + n_exc) + poss.append(pos) + ws.append(ew) + n_exc += pos.numel() + exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) + return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) + + +def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: + """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" + tiles, ks = (n + 15) // 16, k // 16 + p = packed.view(tiles, ks // 4, 3, 32, 16) + sm = ( + p[:, :, :2] + .permute(0, 1, 3, 2, 4) + .reshape(tiles, ks // 4, 32, 4, 8) + .permute(0, 1, 3, 2, 4) + .to(torch.int32) + ) + cw = ( + p[:, :, 2] + .contiguous() + .view(torch.int32) + .view(tiles, ks // 4, 32, 4) + .permute(0, 1, 3, 2) + .long() + & 0xFFFFFFFF + ) + codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( + torch.int32 + ) + bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) + bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) + frag = bits.reshape(tiles, ks, 32, 8) + # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). + w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) + return w.to(torch.int16).view(torch.bfloat16)[:n] + + +def decode( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + n: int, + k: int, + row0: int = 0, + row1: Optional[int] = None, +) -> torch.Tensor: + """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). + ``row0`` must be a multiple of 16.""" + row1 = n if row1 is None else row1 + assert row0 % 16 == 0 + t0, t1 = row0 // 16, (row1 + 15) // 16 + unit_bytes = (k // 64) * 3 * 32 * 16 + dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() + counts = (exc_off[1:] - exc_off[:-1]).long() + tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) + rows = tile_of * 16 + (exc_pos >> 16).long() + cols = (exc_pos & 0xFFFF).long() + sel = (rows >= row0) & (rows < row1) + dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) + return dec + + +def round_trip_ok( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + weight: torch.Tensor, + rows_per_chunk: int = ROWS_PER_CHUNK, +) -> bool: + """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent + window below the BF16 infinity exponent.""" + n, k = weight.shape + if not 0 <= base <= MAX_BASE or k % 64: + return False + if ( + packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 + or exc_off.numel() != (n + 15) // 16 + 1 + ): + return False + for r0 in range(0, n, rows_per_chunk): + r1 = min(n, r0 + rows_per_chunk) + dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) + if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): + return False + return True + + +class M3TargetHead(OpWrapper): + """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit + the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in + another summation order). + + ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` + (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are + the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and + # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with + # a shard whose rows are not a multiple of 16. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{m} verify row(s) of the TP4 target head shard", + spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), + ) + for m in range(1, MAX_ROWS + 1) + ), + Cell( + why="4 rows with outliers and exact zeros in many tiles (the exception lists)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), + ), + Cell( + why="2 rows over a shard of 4000 rows (250 tiles), K 1024", + spec=dict(m=2, n=4000, k=1024, exceptions=True), + ), + ) + + note = """ + The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception + adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs + from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); + integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp + order: the result is deterministic. + """ + + def __call__( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_target_head( + w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl + ) + return out + + def reference( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. + + ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" + weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + return logits.to(torch.bfloat16) + + def is_valid( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> None: + # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count + # stores a partial sum. + assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + assert out.data_ptr() != x.data_ptr(), "out must not alias x" + + +m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/__init__.py new file mode 100644 index 000000000000..77d7b84d9bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py new file mode 100644 index 000000000000..871f6887b7a5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 9..13 decode rows in one CuTe DSL launch (``m3_moe16``): the router, the routed NVFP4 +experts and the shared MXFP8 expert, on caller-owned state (an ``M3Moe16State`` and one ``M3MoeLayer`` per MoE +layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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"] + + +class M3Moe16Engine(OpWrapper): + """One MiniMax-M3 MoE layer for M (9..13 here; the state takes 9..16) decode rows: + ``layer.state.run(layer, hidden, out, arp)``, bf16(shared + routed) [M, 6144]. + + The same arguments and the same math as ``m3_moe_engine`` (the routed part operation for operation, the shared + expert always folded in) on an ``M3Moe16State``: a different schedule for more rows. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # One build per row count (m_max is a trace-time constant); the target routes 9..13 rows here (spec-decode verify + # of 2 x 6 rows lands at 12). + CELLS: tuple[Cell, ...] = ( + Cell(why="9 rows: the smallest build, a partial last slot group", spec=dict(tokens=9)), + Cell(why="10 rows: its own build (m_max is a trace-time constant)", spec=dict(tokens=10)), + Cell(why="11 rows: its own build, an odd row count", spec=dict(tokens=11)), + Cell(why="12 rows: two requests verifying 5-token drafts", spec=dict(tokens=12)), + Cell(why="13 rows: the largest the target routes here", spec=dict(tokens=13)), + ) + + note = """ + Bands as `m3_moe_engine` (same approximations, same quantization boundaries). Not covered: 14..16 rows, which the + state takes and this target never routes here; the all-reduce push (`arp`). No `is_valid`, as `m3_moe_engine`. + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return layer.state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``m3_moe_engine``'s reference with the shared expert; `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=True) + + +m3_moe16_engine = M3Moe16Engine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..f7b53b205fd9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 1..8 decode tokens in one CuTe DSL launch (``m3_moe``): the router, the routed NVFP4 +experts and the shared MXFP8 expert, on caller-owned state (an ``M3MoeState`` and one ``M3MoeLayer`` per MoE layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ...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_107}) + + # One cell per token count: m_max is a trace-time constant, so each is its own build, and the target builds all + # eight (decode batches of 1..8 requests, spec-decode verify rows up to 8). The whys name what the build differs in. + CELLS: tuple[Cell, ...] = ( + Cell( + why="1 token: pre-tiled shared expert, last phase balanced without trailing FC2", + spec=dict(tokens=1), + ), + Cell( + why="2 tokens: pre-tiled shared expert, the one build without AR_PRE", + spec=dict(tokens=2), + ), + Cell( + why="3 tokens: 4 FC2-B builder warps, the shared expert from swizzled scales", + spec=dict(tokens=3), + ), + Cell( + why="4 tokens: one token per routing warp, logits replicated 128x", spec=dict(tokens=4) + ), + Cell( + why="5 tokens: 5 builder warps, a routing warp takes two tokens, logits 32x", + spec=dict(tokens=5), + ), + Cell(why="6 tokens: one request verifying a 5-token draft", spec=dict(tokens=6)), + Cell(why="7 tokens: the B2IMG global FC1 -> FC2 hand-off", spec=dict(tokens=7)), + Cell( + why="8 tokens: B2IMG with a 6-stage ring, 32 distinct experts at most", + spec=dict(tokens=8), + ), + ) + + note = """ + The kernel's elementwise math uses the PTX approximations (rcp.approx, ex2.approx, tanh) where `reference` uses + torch's correctly rounded ones, and its GEMMs sum in tensor-core order where `reference` sums in fp64: a + quantization boundary (NVFP4 activation / intermediate code, routing near-tie) can fall on the other side, so the + two agree within the band the test derives, not bit for bit. The routed part is bit for bit the stock TRTLLM-Gen + NVFP4 MoE (fp4_quantize + fp4_block_scale_moe_runner) on the same routing; the test checks that too. + + Not covered: the all-reduce push (`arp`), which needs a multicast exchange (the decode forward's boundary all-reduce + exercises it); the non-shared state (`shared=False`), which this target never builds. Every layer on one state + must run in one stream order: the workspace is shared and tagged per call, never reset. + + No `is_valid`: the compiled call checks every buffer's dtype, shape and alignment against its signature and + raises on a mismatch (a misaligned `hidden`, an `out` that is not bf16 [M, 6144], a malformed `arp`). + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return layer.state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Plain torch over the layer's buffers (see the class docstring); `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=layer.state.shared) + + +m3_moe_engine = M3MoeEngine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..6e0a44f38b30 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TF32 tensor-core router of the unfused decode MoE (9..16 generation rows), optionally with the NVFP4 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) + + +def nvfp4_quantize( + x: torch.Tensor, sf_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], + the pre-rounding scaled values [M, K]). + + Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. + """ + gs = sf_scale.float().reshape(()) + m, k = x.shape + blocks = x.float().view(m, k // 16, 16) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +def tf32_rna(w: torch.Tensor) -> torch.Tensor: + """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" + bits = w.contiguous().view(torch.int32) + rounded = (bits + 0x1000) & ~0x1FFF + finite = (bits & 0x7F800000) != 0x7F800000 + return torch.where(finite, rounded, bits).view(torch.float32) + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation (``fp4_quantize``'s bytes): returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target routes the unfused MoE's pure generation steps of 9..16 rows here (ROUTER_TC_GEN), logits only. + CELLS: tuple[Cell, ...] = ( + Cell( + why="9 rows: a 3-request verify past the MoE engine", spec=dict(m=9, k=6144, fp4=False) + ), + Cell(why="10 rows", spec=dict(m=10, k=6144, fp4=False)), + Cell(why="11 rows", spec=dict(m=11, k=6144, fp4=False)), + Cell(why="12 rows: 3 requests x 4 tokens", spec=dict(m=12, k=6144, fp4=False)), + Cell(why="13 rows: a partial second token tile", spec=dict(m=13, k=6144, fp4=False)), + Cell(why="14 rows", spec=dict(m=14, k=6144, fp4=False)), + Cell(why="15 rows", spec=dict(m=15, k=6144, fp4=False)), + Cell( + why="16 rows: 4 requests x 4 tokens, two full token tiles", + spec=dict(m=16, k=6144, fp4=False), + ), + ) + + note = """ + The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both + sum the same 16 slice partials in the same order. Only the one the test device picks is driven. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> tuple[torch.Tensor, ...]: + """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" + with true_fp32_matmul(): + logits = x.float() @ tf32_rna(w).t() + if sf_scale is None: + return (logits,) + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return logits, packed, sf + + def is_valid( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + if sf_scale is not None: + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/__init__.py new file mode 100644 index 000000000000..acec8fdbb26f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalization entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..7ba04c92eb61 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step embedding gather and the first layer's input RMSNorm in one CuTe DSL launch, into caller buffers.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm + +__all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] + +VOCAB = 200064 # MiniMax-M3's vocabulary; the decode table is the full (replicated) embedding + + +class M3EmbeddingFirstnorm(OpWrapper): + """Write ``table[ids]`` into ``raw`` (bf16 [M, 6144], a zero row for an id outside ``[0, V)``) and + ``raw * rsqrt(mean(raw^2) + eps) * weight`` into ``out``; returns ``(out, raw)`` -- the first layer's normed input + and its residual. ``ids`` int32 or int64 [M], M = 1..8; ``table`` bf16 [V, 6144] contiguous; ``weight`` bf16 + [6144]. The first call per (M, id dtype, table shape, PDL) compiles (``DecodeEmbedding`` does it at build).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target builds every (rows 1..8, id dtype) owner; the kernel is one CTA per row, so the rows differ only in + # grid size: the ends of the range and both id dtypes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, int32 ids", spec=dict(rows=1, id_dtype=torch.int32)), + Cell(why="eight requests, int32 ids", spec=dict(rows=8, id_dtype=torch.int32)), + Cell(why="one request, int64 ids", spec=dict(rows=1, id_dtype=torch.int64)), + Cell( + why="six rows (one request verifying a draft), int64 ids", + spec=dict(rows=6, id_dtype=torch.int64), + ), + # Every other row count the decode dispatch takes (modeling.EMBED_FIRSTNORM), both id dtypes. + *( + Cell( + why=f"{rows} rows: a decode row count between the ones above", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in (2, 3, 4, 5, 7) + ), + ) + + note = """ + The gathered rows are exact. The norm is flashinfer's RMSNorm arithmetic in its reduction order with a fast + rsqrt, where `reference` sums in fp64 and takes torch's rsqrt: within the band the test derives, not bit for bit. + An id outside [0, V) gives a zero row and a zero norm (torch's embedding would raise). No accepted input is known + to compute wrongly: the host side rejects any other dtype, shape or a strided buffer. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return m3_embed_norm(ids, table, weight, eps, raw, out, use_pdl) + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x * rstd * weight.double()).to(torch.bfloat16), rows + + +m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/__init__.py new file mode 100644 index 000000000000..d25c41dc9557 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries this target alone calls: the Eagle3 verify ops, the capture / fc_norm RMSNorms and +the step-input staging.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..be5aa46fe774 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_chunk_norm import rmsnorm_rows + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = ( + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ) + + note = """ + fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional + one-ulp move against torch's statistic. + """ + + def __call__( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) + return dst + + def reference( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" + x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h + return rmsnorm_rows(x, w, eps) + + def is_valid( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> None: + # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read + # past its elements without an error. + assert w.is_contiguous(), "w must be dense [H]" + + +m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..6699a55fc3e6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states +(``m3_spec_tail.chunk_norm``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ) + + note = """ + fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then + against torch's statistic; the cells' band states it. + """ + + def __call__( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) + return dst + + def reference( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" + c, h = w.shape + x = src[:, : c * h].reshape(src.shape[0], c, h) + return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) + + def is_valid( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ) -> None: + c, h = w.shape + assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + + +m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..8f0f69ee812b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat +(``m3_spec_tail.draft_input``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3DraftInput(OpWrapper): + """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = + hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of + a vocab-local embedding). Returns ``(out, residual)``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling.DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the + # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests + # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", + spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", + spec=dict(m=n, h=6144, gather=True, invalid_id=False), + ) + for n in range(1, 9) + ), + Cell( + why="third pass of 3 requests: one row each, hidden read in place", + spec=dict(m=3, h=6144, gather=False, invalid_id=False), + ), + Cell( + why="one row with an out-of-table id: the embedding half is zero", + spec=dict(m=1, h=6144, gather=False, invalid_id=True), + ), + ) + + note = """ + The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's + summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, + also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the + drafter's embedding is the whole table, so its ids are always in range. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl + ) + return out, residual + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / + ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" + valid = (ids >= 0) & (ids < table.shape[0]) + emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) + emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) + hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] + ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) + return ref, (hid.clone() if residual is not None else None) + + def is_valid( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> None: + h = table.shape[1] + # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. + assert table.is_contiguous() and hidden.is_contiguous(), ( + "table / hidden must be dense [*, H]" + ) + if rows is not None: + assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" + else: + assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" + # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + assert residual is None or ( + residual.is_contiguous() and residual.shape == (ids.shape[0], h) + ), "residual must be contiguous [M, H]" + + +m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py new file mode 100644 index 000000000000..75626e570244 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter MLP's plain SwiGLU as one C++ launch (``modeling.DRAFT_SILU_CPP``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3DraftSilu(OpWrapper): + """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: + trtllm::minimax_m3_sm107_tp4_silu_and_mul (minimaxM3Sm107Tp4/siluAndMul.cu), bit for bit the Triton + ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP + runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs + the Triton op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP4 drafter MLP shard: intermediate 18432 / 4 = 4608, gate | up 9216 wide, at the drafter's decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP4 drafter MLP at 1 row", spec=dict(rows=1, d=4608)), + Cell(why="the TP4 drafter MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=4608)), + Cell(why="the TP4 drafter MLP at 8 rows", spec=dict(rows=8, d=4608)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (the same PTX instruction sequence), checked over every + bf16 gate and every bf16 up pattern. + """ + + def __call__(self, x: torch.Tensor) -> Optional[torch.Tensor]: + if ( + x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 16 != 0 + or x.data_ptr() % 16 != 0 + or x.shape[0] > 65535 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp4_silu_and_mul(x, out) + return out + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul(x) + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_draft_silu = M3DraftSilu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..5dbb324a7d50 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One pass of the fused greedy Eagle3 draft loop: the drafter's vocab-sharded BF16 head (streamed from its lossless +12-bit coding, ``ops.gemm.m3_target_head``, or as BF16 fragments), the cross-rank argmax and the loop's per-pass +bookkeeping in one kernel (``m3_draft_step.draft_step``).""" + +from __future__ import annotations + +import functools +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ..gemm.m3_target_head import decode as decode_coded +from .m3_shard_accept import exchange_tokens, local_argmax, make_keys + +#: Requests (head rows) per pass, and exchange rows per call half of the workspace. +MAX_TOKENS = 8 +#: 16-row weight tiles one CTA can own (the kernel's shared-memory accumulator). +MAX_TILES_PER_CTA = 32 +KERNEL_WARPS = 16 # warps per CTA (draftStep.h kWarps) + + +@functools.lru_cache(maxsize=64) +def max_tile_pieces(n_tiles: int, num_ctas: int, chunks_per_tile: int) -> int: + """The most warps that stream one weight tile under the kernel's whole-tile split: CTA c owns tiles + [c * n / grid, (c + 1) * n / grid), and its (tile, chunk) range is cut into KERNEL_WARPS equal contiguous + segments. The shared-memory atomics add a tile's pieces in arrival order: deterministic for at most 2 pieces (two + addends from 0 commute), not for more (fixed_order then sums them in warp order).""" + worst = 1 + for c in range(num_ctas): + lo, hi = c * n_tiles // num_ctas, (c + 1) * n_tiles // num_ctas + span = (hi - lo) * chunks_per_tile + bounds = [w * span // KERNEL_WARPS for w in range(KERNEL_WARPS + 1)] + for t in range(hi - lo): + a, b = t * chunks_per_tile, (t + 1) * chunks_per_tile + worst = max( + worst, sum(1 for w in range(KERNEL_WARPS) if bounds[w] < b and bounds[w + 1] > a) + ) + return worst + + +#: Draft length of the fused loop, and the drafter's TP4 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 50016 +HIDDEN = 6144 + + +def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the + 12-bit coded stream with its exception lists decoded.""" + if exc_off is None: + return from_fragments(wf) + n = (exc_off.numel() - 1) * 16 + return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + + +def shard_logits( + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() + + +class M3DraftStep(OpWrapper): + """One fused greedy draft pass over ``rows <= 8`` requests (one head row each). + + Head input ``x[x_rows[b]]`` (pass 0: the last accepted verify row of each request) or ``x[b]``. The head is this + rank's BF16 lm_head shard in fragment order (``m3_draft_step.bf16_fragments``, global ids ``vocab_offset ..``), + or with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base`` its lossless 12-bit coding (``wf`` = the uint8 stream + of ``m3_target_head.encode_chunked``, which must round-trip to the BF16 shard); each logit is the fp32 dot product + rounded to BF16 and the row's token is the global argmax over all ranks' + shards (ties, -0 / +0 included, to the lowest global id), exchanged through the multicast slots + ``[2][8][world]``. The last CTA then writes ``draft_out[b, step]`` and ``ids_next[b]`` (the token), ``pos_next[b] + = pos_src[x_rows[b] or b] + 1``, and advances the draft KV lengths as the in-tree draft loop does between passes + (pass 0 with ``num_acc``: ``kv_lens -= draft_len - num_acc`` and ``seq_lens = 1``; later passes: ``kv_lens += + 1``). ``out_keys`` (optional) receives each row's winning key. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling.DRAFT_STEP): draft_len 3 at 1..8 requests, the drafter's + # full-vocabulary BF16 head (50016 rows per rank of 4, hidden 6144). Pass 0 reads the verify rows' hidden (4 per + # request) through x_rows and rolls the KV back; the later passes read one row per request. + _PASSES: tuple[Cell, ...] = ( + *( + Cell( + why=f"pass 0 at {n} request(s): head rows gathered from {4 * n} verify rows, KV rollback, seq_lens = 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=0, + rank=n % 4, + world=4, + vs=DRAFT_SHARD, + k=HIDDEN, + ), + ) + for n in range(1, MAX_TOKENS + 1) + ), + Cell( + why="pass 1 at two requests: one row each, kv_lens + 1, rank 3's ids", + spec=dict( + nreq=2, draft_len=DRAFT_LEN, step=1, rank=3, world=4, vs=DRAFT_SHARD, k=HIDDEN + ), + ), + Cell( + why="last pass (step 2) at 8 requests, rank 0", + spec=dict( + nreq=8, draft_len=DRAFT_LEN, step=2, rank=0, world=4, vs=DRAFT_SHARD, k=HIDDEN + ), + ), + Cell( + why="pass 2 at 5 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=5, + draft_len=DRAFT_LEN, + step=2, + rank=1, + world=4, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + ), + ), + ) + # The coded head (the fused draft loop's default) at every pass shape, and the BF16 fragments (the coding's + # fail-closed path) at the first and last pass shapes and with ties. + CELLS: tuple[Cell, ...] = ( + *(Cell(why=f"coded: {c.why}", spec=dict(c.spec, coded=True)) for c in _PASSES), + *( + Cell(why=f"bf16: {c.why}", spec=dict(c.spec, coded=False)) + for c in (_PASSES[0], _PASSES[MAX_TOKENS - 1], _PASSES[-1]) + ), + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit; + the coded head's weight is the BF16 shard bit for bit; its exceptions decode to +0 in the MMAs and their exact + products are added at the tile's flush) + and the argmax is the in-tree one over BF16 logits, so a token differs from the in-tree path's only where the fp32 + sum, accumulated in another order, rounds to the other side of a BF16 boundary and that decides the row's maximum. + The cells' rows are decidable and the peers' keys sit 0.25 away (or tie exactly); the edge tests (forced BF16 + ties, -0 / +0, top-2 one fp32 step apart that round to one BF16 value, adjacent BF16 values) use integer data + whose fp32 sums are exact in any order and hold the kernel to the in-tree path bit for bit. The cross-rank + exchange is certified on one device with the peers' keys written into the slots (see m3_shard_accept). + """ + + def __call__( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> None: + # fold: (the 14 fold-tail tensors in buildParams' order, rows, tokens per step) -- the step's forward-tail store + # and next-step staging, run by the pass's last CTA after its tokens (unused by this target, which passes None). + tensors, fold_rows, fold_t = ( + (list(fold[0]), int(fold[1]), int(fold[2])) if fold is not None else (None, 0, 0) + ) + torch.ops.trtllm.minimax_m3_sm107tp4_draft_step( + x, + x_rows, + rows, + wf, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + num_ctas, + pdl, + out_keys, + exc_off, + exc_pos, + exc_w, + exp_base, + tensors, + fold_rows, + fold_t, + fixed_order, + ) + + def reference( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> dict[str, torch.Tensor]: + """Every buffer the pass writes, computed on copies: "logits" (this rank's BF16 logits as fp32), "local_keys" + (this rank's key per row), "keys" (the winning key), "tokens", "draft_out", "ids_next", "pos_next", "kv_lens" + and "seq_lens". + + Reads ``ws_uc``, ``counters`` and ``kv_lens`` before the call. ``keys``, ``ws_mc_ptr``, ``num_ctas``, ``pdl``, + ``out_keys`` and ``fixed_order`` do not change the result; ``fold`` adds the fold tail's writes (m3_fold_tail), + not these. + """ + h = x[x_rows.long()] if x_rows is not None else x[:rows] + logits = shard_logits(h, wf, exc_off, exc_pos, exc_w, exp_base) + best, idx = local_argmax(logits, vocab_offset) + local_keys = make_keys(best, idx) + tokens = exchange_tokens(local_keys, ws_uc, counters, rank, world, MAX_TOKENS) + half = int(counters[1].item()) & 1 + slots = ws_uc.view(torch.int64)[ + half * MAX_TOKENS * world : (half * MAX_TOKENS + rows) * world + ] + slots = slots.view(rows, world).clone() + slots[:, rank] = local_keys + sign = torch.iinfo(torch.int64).min + win = (slots ^ sign).max(dim=1).values ^ sign + d_out = draft_out.clone() + d_out[:rows, step] = tokens + src = x_rows.long() if x_rows is not None else torch.arange(rows, device=x.device) + lens = kv_lens.clone() + sl = seq_lens.clone() if seq_lens is not None else None + if step == 0 and num_acc is not None: + lens[:rows] -= draft_len - num_acc[:rows] + sl[:rows] = 1 + elif step > 0: + lens[:rows] += 1 + out = { + "logits": logits, + "local_keys": local_keys, + "keys": win, + "tokens": tokens, + "draft_out": d_out, + "ids_next": tokens.clone(), + "pos_next": pos_src[src] + 1, + "kv_lens": lens, + } + if sl is not None: + out["seq_lens"] = sl + return out + + def is_valid( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> None: + assert not bool(keys[:rows].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + # The kernel reads the head row x + x_rows[b] * K and the position pos_src[x_rows[b]] with no bound. + if x_rows is not None: + assert bool(((x_rows[:rows] >= 0) & (x_rows[:rows] < x.shape[0])).all()), ( + "x_rows must index x" + ) + assert bool((x_rows[:rows] < pos_src.shape[0]).all()), ( + "pos_src must cover the head rows" + ) + else: + assert pos_src.shape[0] >= rows, "pos_src must cover the head rows" + # Each CTA accumulates at most MAX_TILES_PER_CTA tiles in shared memory. + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + assert -(-tiles // num_ctas) <= MAX_TILES_PER_CTA, "too many weight tiles per CTA" + # The atomics' arrival-order sum is deterministic only while no tile takes more than 2 warp pieces. + assert ( + fixed_order + or max_tile_pieces(tiles, num_ctas, x.shape[1] // (16 if exc_off is None else 64)) <= 2 + ), ( + "a tile takes more than 2 warp pieces on this grid: the atomics' sum is order-dependent (fixed_order)" + ) + if exc_off is not None: + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), ( + "exc_off must end at the exception count" + ) + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + # A NaN logit maps to a key above every finite one and wins the row silently. + h = x[x_rows.long()] if x_rows is not None else x[:rows] + assert not bool(torch.isnan(h).any()), "NaN head inputs win the argmax" + if fold is not None: + assert len(fold) == 3 and len(fold[0]) == 14, ( + "fold is (the 14 fold-tail tensors, rows, tokens per step)" + ) + if num_acc is not None: + na = num_acc[:rows] + assert bool(((na >= 1) & (na <= draft_len + 1)).all()), ( + "num_acc counts must lie in [1, draft_len + 1]" + ) + + +m3_draft_step = M3DraftStep() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..fcb327e374e3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + force_vec_kernel[(1,)]( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + """The counts (a new tensor; ``counter`` and ``out`` are not written).""" + c = int(counter.item()) + 1 + slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) + idx = (c * counter_stride + slots * slot_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec = M3ForceVec() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py new file mode 100644 index 000000000000..5a2ffd8b6627 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, each request drawn at its own key.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.m3_force_vec import force_vec_req_kernel + + +class M3ForceVecReq(OpWrapper): + """For each of the ``batch_size`` requests ``b`` write ``min(base_total + (pool[(keys[b] * key_stride) & + pool_mask] < frac), max_total)`` into ``out`` (int32): the forced-acceptance draw from the fixed-seed fp32 ``pool`` + at the request's own key (int32 ``keys``, its KV length at this verify), so a request's draws depend on neither the + server's history nor the other requests. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or + 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it, with ``FORCE_AL_PER_REQUEST``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + force_vec_req_kernel[(1,)]( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + """The counts (a new tensor; ``out`` is not written).""" + idx = (keys[:batch_size].to(torch.int64) * key_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec_req = M3ForceVecReq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..7623f5fe582b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampled (temperature + top-p) Eagle3 verification over the TP-sharded vocabulary in one kernel +(``m3_sampled_accept.sampled_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_shard_accept import strict_accept + +#: Verify rows per call. +MAX_ROWS = 48 +#: The fused sampling kernel's 1 / T guard. +TEMP_EPSILON = 1e-6 + + +def nucleus_probs( + logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +class M3SampledAccept(OpWrapper): + """Sampled verification of ``rows = nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard. + + Each row draws one token from its top-p nucleus of softmax(logits / T) over the full (all-rank) vocabulary by + distributed rejection sampling (Philox keyed by the row's seed / offset, the row, the attempt and the CTA), then + the strict acceptance runs against ``draft`` as in ``m3_shard_accept``: ``acc_tokens``, ``num_acc``, + ``gather_ids``, optional ``next_new[:, 0]``. ``tokens`` receives the sampled token per row; ``extra_attempts`` + accumulates exchange rounds beyond the second. The exchange goes through the multicast workspace (two call + halves of ``minimax_m3_sm107tp4_sampled_accept_ws_words(world)`` words) with the (finished rows, call) ``counters``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3 (T = 4 rows per request) at the TP4 caps: 1 / 2 / 8 requests (4 / 8 / 32 rows, 8 CTAs per row) + # and the 48-row maximum (4 CTAs per row); the 50016-column target shard. Single rank (see note). + CELLS: tuple[Cell, ...] = ( + Cell( + why="cap 1 (4 rows), top_p 0.8 over a 12-token head: the draws follow the nucleus distribution", + spec=dict( + nreq=1, vs=50016, vocab_offset=0, check="distribution", top_p=0.8, temperature=0.7 + ), + ), + Cell( + why="cap 8 (32 rows, 256 co-resident CTAs): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=8, + vs=50016, + vocab_offset=2 * 50016, + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ), + Cell( + why="cap 2 (8 rows) with top_p = 1 (no cut: the first draw stands), next_new and forced counts", + spec=dict( + nreq=2, + vs=50016, + vocab_offset=50016, + check="distribution", + top_p=1.0, + temperature=1.0, + force="vec", + ), + ), + Cell( + why="48 rows (4 CTAs per row): the second launch geometry, nucleus distribution at top_p 0.5", + spec=dict( + nreq=12, + vs=50016, + vocab_offset=3 * 50016, + check="distribution", + top_p=0.5, + temperature=1.3, + force="total", + ), + ), + ) + + note = """ + What is sampled is gated statistically, the rest exactly. Every drawn token must lie in the reference nucleus; + the empirical distribution over many calls (distinct Philox offsets) must sit within a stated total-variation + band of the reference's; the acceptance outputs must equal the strict acceptance of the drawn tokens exactly. + The Philox stream itself is not reproduced in torch: the kernel's draw uses the hardware exp / log + approximations (__expf, __logf) and a CTA-order inverse CDF, which torch does not reproduce bit for bit. + + Certified at world 1 only. The decision is replicated from every rank's per-round records, which another rank + computes from its own shard on the fly, so peers cannot be stood in for on one device; the cross-rank path + needs the multi-device job and is not covered here. Rows with top-k or min-p must not reach this op (it applies + neither); is_valid states it. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + prewait_params: bool = True, + early_trigger: bool = False, + one_round_k: int = 0, + decided: Optional[torch.Tensor] = None, + ilp_walk: bool = False, + ) -> None: + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept( + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + prewait_params, + early_trigger, + one_round_k, + decided, + ilp_walk, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + prewait_params: bool = True, + early_trigger: bool = False, + one_round_k: int = 0, + decided: Optional[torch.Tensor] = None, + ilp_walk: bool = False, + ) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + prewait_params: bool = True, + early_trigger: bool = False, + one_round_k: int = 0, + decided: Optional[torch.Tensor] = None, + ilp_walk: bool = False, + ) -> None: + rows = logits.shape[0] + # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the + # wrong rows. + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + # Per-row parameters are read at the row index, unchecked. + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), ( + "per-row parameters must cover every row" + ) + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + # The finished-row counter elects the CTA that runs the acceptance. + assert int(counters[0].item()) == 0, ( + "counters[0] (finished rows) must be zero between calls" + ) + + +m3_sampled_accept = M3SampledAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..9193b3dce61d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept_ws_words.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Size of the sampled-verify exchange workspace (``m3_sampled_accept.SampledAcceptWorkspace.create``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Exchange rounds per row, verify rows per call, CTAs per row (at most), uint64 words per record. +MAX_ATTEMPTS = 16 +MAX_ROWS = 48 +MAX_SPLITS = 8 +RECORD_WORDS = 4 +#: SA-1R (one_round_k 9 / 82): 9 listed bf16 values ride each record, 3 per extra word. +LIST_WORDS = {0: 0, 9: 3, 82: 3} + + +class M3SampledAcceptWsWords(OpWrapper): + """uint64 words of the multicast workspace ``minimax_m3_sm107tp4_sampled_accept`` exchanges through for + ``world`` ranks: two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="the tp4 target's workspace", spec=dict(world=4)), + Cell(why="a single rank (local workspace, no multicast)", spec=dict(world=1)), + Cell( + why="the tp4 target's workspace with SA-1R pooled (82)", + spec=dict(world=4, one_round_k=82), + ), + Cell( + why="the tp4 target's workspace with SA-1R top-9 (9)", spec=dict(world=4, one_round_k=9) + ), + ) + + note = """ + A host computation (no device work); it is an op so the workspace size and the kernel's indexing come from one + definition. + """ + + def __call__(self, world: int, one_round_k: int = 0) -> int: + return torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept_ws_words(world, one_round_k) + + def reference(self, world: int, one_round_k: int = 0) -> int: + """Two halves of attempts x rows x (world x splits) records x record words.""" + return ( + 2 + * MAX_ATTEMPTS + * MAX_ROWS + * world + * MAX_SPLITS + * (RECORD_WORDS + LIST_WORDS[one_round_k]) + ) + + +m3_sampled_accept_ws_words = M3SampledAcceptWsWords() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..8beba4373b02 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Greedy Eagle3 verification from this rank's logits shard: per-row argmax, the cross-rank (value, index) exchange +and the strict acceptance in one kernel (``m3_shard_accept.shard_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Verify rows per call, and exchange rows per call half of its workspace. +SHARD_ROWS = 48 + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out + + +class M3ShardAccept(OpWrapper): + """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` + (global ids ``vocab_offset ..``). + + Each row's token is its global argmax over all ranks' shards: keys ``ordered(value) << 32 | ~index`` (ties take + the lowest index, as ``torch.argmax``) are max-reduced over the shard, pushed into every rank's exchange slot + through the multicast slots ``[2][48][world]`` (Lamport: 0 = empty) and max-reduced over the ranks by the last + CTA, which then runs the strict acceptance: ``acc_tokens [nreq, D + 1]``, ``num_acc [nreq]`` (``force_vec`` / + ``force_total`` override the counts), ``gather_ids [nreq]`` (the first draft pass's row per request) and, when + given, ``next_new[:, 0]`` (the last accepted token). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every greedy verify step of the target (modeling.SHARD_ACCEPT): draft_len 3 at 1..12 requests (4..48 rows), the + # target lm_head's BF16 shard logits (50016 columns per rank at TP4). The logits come from the head GEMM, so the + # op only reduces. + CELLS: tuple[Cell, ...] = ( + Cell( + why="cap 1: 4 verify rows (1 request x T = 4), rank 0 of 4, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=50016, world=4, rank=0, force="none", strided=False), + ), + Cell( + why="cap 2: 8 verify rows, rank 2, global ids offset by 2 x 50016", + spec=dict(nreq=2, draft_len=3, vs=50016, world=4, rank=2, force="none", strided=False), + ), + Cell( + why="cap 4: 16 verify rows (4 requests x T = 4), rank 1 of 4, next_new staged", + spec=dict(nreq=4, draft_len=3, vs=50016, world=4, rank=1, force="none", strided=False), + ), + Cell( + why="cap 8: 32 verify rows, rank 3, global ids offset by 3 x 50016", + spec=dict(nreq=8, draft_len=3, vs=50016, world=4, rank=3, force="none", strided=False), + ), + Cell( + why="12 rows (3 requests) read from a row-strided view of a wider logits buffer, forced counts", + spec=dict(nreq=3, draft_len=3, vs=50016, world=4, rank=0, force="vec", strided=True), + ), + Cell( + why="the 48-row maximum (12 requests at draft 3), force_total, rows with exact cross-rank ties", + spec=dict( + nreq=12, + draft_len=3, + vs=50016, + world=4, + rank=2, + force="total", + strided=False, + ties=True, + ), + ), + # Every other request count of the verify range, on the 4-rank 50016-column shards. + *( + Cell( + why=f"{n} requests ({4 * n} verify rows): a request count between the ones above", + spec=dict( + nreq=n, draft_len=3, vs=50016, world=4, rank=n % 4, force="none", strided=False + ), + ) + for n in (5, 6, 7, 9, 10, 11) + ), + ) + + note = """ + The op reduces BF16 logits exactly, so the reference is bit-exact, including cross-rank ties (lowest global + index wins). The cross-rank exchange is certified on one device: the other ranks' keys are written into this + rank's slots before the call, which is what their multicast pushes would have put there; the multicast mapping + itself (ws_mc_ptr = this rank's slots, as at world 1) is not exercised across devices here. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + num_ctas, + force_vec, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Lowest-index maximum of each row's BF16 logits, the largest key among it and the peers' keys in + ``ws_uc``, then the strict acceptance. Reads ``ws_uc`` / ``counters`` before the call; the scratch and + output buffers and ``num_ctas`` do not change the result.""" + best, idx = local_argmax(logits.float(), vocab_offset) + tokens = exchange_tokens(make_keys(best, idx), ws_uc, counters, rank, world, SHARD_ROWS) + out = strict_accept(tokens, draft, batch_idx, force_total, next_new, force_vec) + out["tokens"] = tokens + return out + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[: logits.shape[0]].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + assert ws_uc.numel() >= 2 * SHARD_ROWS * world, "ws_uc must hold [2][48][world] slots" + # A NaN logit maps to a key above every finite one and wins the row silently. + assert not bool(torch.isnan(logits).any()), "NaN logits win the argmax" + nreq = draft.shape[0] + assert bool(((batch_idx[:nreq] >= 0) & (batch_idx[:nreq] < nreq)).all()), ( + "batch_idx must index the nreq requests" + ) + if force_vec is not None: + fv = force_vec[:nreq] + assert bool(((fv >= 1) & (fv <= draft.shape[1] + 1)).all()), ( + "force_vec counts must lie in [1, D + 1]" + ) + + +m3_shard_accept = M3ShardAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..ff07c54de1f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> +``m3_spec_tail.slot_scatter``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: + """src rows zero-padded (or cut) to ``width`` columns.""" + out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) + w = min(width, src.shape[1]) + out[:, :w] = src[:, :w] + return out + + +class M3SlotScatter(OpWrapper): + """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, + ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The + destinations are the store's transposed [slots, width] views.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of + # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) + # is zero-padded. + CELLS: tuple[Cell, ...] = ( + Cell( + why="8 rows into shuffled slots, widths equal to the store's", + spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), + ), + Cell( + why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", + spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), + ), + ) + + note = """ + Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. + """ + + def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): + torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) + + def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): + """The four destinations after the call, on copies.""" + s = slots.long() + out = {} + for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): + r = dst.clone() + r[s] = padded_rows(src[: s.numel()], dst.shape[1]) + out[name] = r + lr = dl.clone() + lr[s] = sl[: s.numel()] + out["dl"] = lr + return out + + def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): + assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( + "slots must index the store" + ) + # Two rows into one slot race (which one lands is unordered). + assert slots.unique().numel() == slots.numel(), "slots must be distinct" + assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + + +m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py new file mode 100644 index 000000000000..36e421412f30 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's speculative step correction as one launch (``RuntimeHooks.spec_offsets_op``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger + +# The engagement line is logged once per process (module state: __call__ reads no instance state). +_ENGAGED: list = [] + + +class M3SpecOffsets(OpWrapper): + """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp4_spec_offsets (minimaxM3Sm107Tp4/specOffsets.cu), the + same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other + layouts (the caller runs the ATen adds).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="1 request, 4 tokens (draft length 3)", spec=dict(tokens=4, seqs=1)), + Cell(why="8 requests, 32 tokens", spec=dict(tokens=32, seqs=8)), + Cell(why="wrap-around values", spec=dict(tokens=8, seqs=2, wrap=True)), + ) + + note = "Integer adds: the reference (the two in-place ATen adds) is exact." + + def __call__(self, pos, pos_off, kv, kv_off) -> bool: + ts = (pos, pos_off, kv, kv_off) + if ( + any( + t.dtype != torch.int32 or t.dim() != 1 or t.stride(0) != 1 or not t.is_cuda + for t in ts + ) + or pos.numel() != pos_off.numel() + or kv.numel() != kv_off.numel() + ): + return False + torch.ops.trtllm.minimax_m3_sm107_tp4_spec_offsets(pos, pos_off, kv, kv_off) + if not _ENGAGED: + _ENGAGED.append(True) + logger.info("[M3 host] spec step correction: one launch taken (spec_offsets_op)") + return True + + def reference(self, pos, pos_off, kv, kv_off) -> None: + pos += pos_off + kv += kv_off + + def is_valid(self, pos, pos_off, kv, kv_off) -> None: + assert pos.numel() == pos_off.numel() and kv.numel() == kv_off.numel(), ( + "equal lengths per pair" + ) + + +m3_spec_offsets = M3SpecOffsets() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..16c4000199da --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch +(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS: tuple[Cell, ...] = ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + ) + + note = """ + Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid + states them. + """ + + def __call__( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ) + + def reference( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" + nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) + s = slots[:P].long() + ids_r, draft_r = ids.clone(), draft.clone() + ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) + draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) + pos_r = torch.zeros_like(pos_off) + pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] + kv_r = torch.zeros_like(kv_off) + kv_r[b0 : b0 + P] = lens[s] - T + return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} + + def is_valid( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or + # writes another buffer. + s = slots[:P] + assert bool( + ( + (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) + ).all() + ), "slots must index the store" + assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( + "ppi must index lens" + ) + assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" + assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" + assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( + "offset buffers too short" + ) + + +m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py new file mode 100644 index 000000000000..122d43fb175a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode graph's step-input prologue in one launch (this target's ``RuntimeHooks.pack_prologue_op``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepPrologue(OpWrapper): + """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch + (trtllm::minimax_m3_sm107_tp4_step_prologue, minimaxM3Sm107Tp4/stepPrologue.cu). + + ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one + ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the + window, a device copy's a device address). The launch writes up to ``max_desc`` copies of each kind, every host + copy before any device copy, then zeroes the header; with ``epoch`` it increments ``epoch[0]`` once. The default + prologue does the same with two parallel-scatter launches (``step_inputs._graph_scatter_kernel``) and a header + fill.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP4 decode step's packs: a few dozen host copies of 4 to a few hundred bytes (positions, KV lengths, slots, + # the spec sampler's uploads), a handful of device copies (the graph runner's static inputs), some reading what a + # host copy of the same window wrote. + CELLS: tuple[Cell, ...] = ( + Cell(why="an empty window (a replay with no new upload)", spec=dict(h2d=0, d2d=0)), + Cell(why="host copies only, odd sizes and alignments", spec=dict(h2d=40, d2d=0)), + Cell( + why="host copies feeding device copies (a device copy reads a host copy's destination)", + spec=dict(h2d=24, d2d=8, chained=True), + ), + Cell( + why="every descriptor slot (128 + 128 of the pack's 256)", spec=dict(h2d=128, d2d=128) + ), + ) + + note = """ + Byte copies: the reference (the default prologue's Triton launches and fill on the same window) is exact. + Destinations are raw addresses; is_valid cannot check them. + """ + + def __call__(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + torch.ops.trtllm.minimax_m3_sm107_tp4_step_prologue(window, epoch, max_desc) + + def reference(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + """The default prologue on the same arguments: it writes the same destinations.""" + from tensorrt_llm._torch.pyexecutor.step_inputs import _graph_scatter_kernel + + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + bump = epoch is not None and not device_src + _graph_scatter_kernel[(max_desc,)]( + window, + stage_i64, + epoch if bump else stage_i64, + DEVICE_SRC=device_src, + BLOCK=256, + BUMP_EPOCH=bump, + num_warps=1, + ) + window[:16].zero_() + + def is_valid(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + assert window.dtype == torch.uint8 and window.is_contiguous(), "window: contiguous uint8" + assert window.numel() >= 16 + 24 * max_desc, "window: header and max_desc descriptors" + counts = window[:16].view(torch.int64).tolist() + assert counts[0] + counts[1] <= max_desc, "at most max_desc copies in all" + + +m3_step_prologue = M3StepPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/weights.py new file mode 100644 index 000000000000..07fbe7f06468 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/weights.py @@ -0,0 +1,865 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_107 / tp4. + +Two tables, one per checkpoint the target reads: + +* ``MODEL_WEIGHTS`` (``MiniMaxM3Tp4Weights``): the vision-language checkpoint's language model, + ``language_model.*``. Every rank is handed the whole checkpoint, so every tensor-parallel split is this + table's: the Q heads, the K/V/index-Q heads (one each per rank), the dense MLPs and the shared expert's + intermediate, the o_proj / down-projection input columns, every routed expert's intermediate (all 128 + experts live on every rank), and the vocabulary of the shell's ``lm_head``. The index-K head is replicated. +* ``DRAFT_WEIGHTS`` (``MiniMaxM3Tp4DraftWeights``): the one-layer BF16 Eagle3 drafter's own checkpoint + (``speculative_model``), loaded through the shell's ``load_draft_weights`` when speculation is on. + +Storage is HF ``[out, in]`` row-major. Five families need a transform, each a layout a kernel demands: + +* **MXFP8 block scales** (every attention projection, the dense MLPs and the shared expert): the checkpoint + stores UE8M0 bytes ``[N, K / 32]``; the GEMMs read the 128x4-swizzled order + (``trtllm::block_scale_interleave``). A fused projection concatenates its shards' scales first and swizzles + once. +* **The five-way sparse-layer projection** ``[Q | K | V | index-Q | index-K]`` and the three-way dense one + ``[Q | K | V]``: the per-rank shards concatenated along N. +* **Gated MLPs** (dense and shared expert): ``[gate ; up]`` rows. +* **NVFP4 routed experts** (layers 3..59): ``[up (w3) ; gate (w1)]`` rows, then the trtllm-gen row shuffle of + the weight and of its E4M3 scales (gated-activation reorder + 128-row epilogue tile shuffle, + ``trtllmgen_maybe_get_cached_*_permute_indices``) and the 128x4 swizzle of the scales. +* **Gemma boundary norms**: ``input_layernorm``, ``post_attention_layernorm`` and ``model.norm`` are stored as + Gemma norms, ``(1 + w) * x``; the fused all-reduce + RMSNorm kernels apply ``w * x``, so ``1`` is added + (in BF16) at load. The per-head Q/K/index norms stay Gemma: their kernels apply the offset themselves. + +The routed experts' per-tensor scalars (input scales and second-level weight scales, every expert) load raw; +``modeling`` derives the kernels' global scales, alphas and activation constants from them once the weights +are real. + +Loading contract: ``load(model, weights)`` consumes the engine-provided checkpoint dict, fills every declared +parameter exactly once and asserts coverage in both directions: every declared parameter written, and every +checkpoint key either consumed or in ``expected_unconsumed`` (the vision tower and its projector, which this +text-only target does not build, and the language model's multi-token-prediction head). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) + +#: ``layers=ALL_LAYERS`` is the common case: a weight every decoder layer carries. +ALL_LAYERS = object() + +_LM = "language_model." +#: Checkpoint key families this target does not build: the vision tower and its two projectors (text only), and the +#: language model's multi-token-prediction head (speculation, when on, is the separate Eagle3 drafter). +_UNBUILT_PREFIXES = ( + "vision_tower.", + "multi_modal_projector.", + "patch_merge_mlp.", + _LM + "model.mtp.", +) + +MX_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 elements along K +NV_BLOCK = 16 # NVFP4: one E4M3 scale per 16 elements along K +EPILOGUE_TILE_M = ( + 128 # the trtllm-gen MoE runner's row tile: fixes its weight and scale row shuffles +) + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + ``shape`` takes the dimension bundle and returns the declared shape. ``src`` is a callable + ``(d, i) -> [(key or key tuple, index into the param | None), ...]`` with keys relative to the checkpoint + root; a key tuple hands the transform those tensors together. ``layers`` is ``ALL_LAYERS``, ``None`` (not + per-layer) or ``d -> layer indices``. ``transform(src, d)`` runs on the destination's device before the copy. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + +# ------------------------------------------------------------------------------------------------ shards + + +def _col( + t: torch.Tensor, d: SimpleNamespace, tp: int | None = None, rank: int | None = None +) -> torch.Tensor: + """This rank's rows (column-parallel output split), even split, sliced on the host and moved to the + destination device (only this rank's bytes are read).""" + tp = d.tp if tp is None else tp + rank = d.rank if rank is None else rank + n = t.shape[0] // tp + return _dev(t[rank * n : (rank + 1) * n], d) + + +def _row(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """This rank's input columns (row-parallel input split), even split, on the destination device.""" + n = t.shape[1] // d.tp + return _dev(t[:, d.rank * n : (d.rank + 1) * n], d) + + +def _dev(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """``t`` on the destination device (``d.dev``, set per fill).""" + return t.to(d.dev, non_blocking=True) + + +def _swizzle(scale: torch.Tensor) -> torch.Tensor: + """128x4 swizzle of a row-major ``[N, K / 32]`` UE8M0 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(scale.contiguous()) + + +# --------------------------------------------------------------------------------------------- transforms + + +def _gemma_fold(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A Gemma boundary norm as the plain norm the fused all-reduce epilogue applies: ``w + 1`` in BF16.""" + return t + 1.0 + + +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) + + +def _qkv_sparse_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[Q | K | V | index-Q | index-K]``: Q by attention heads, K/V/index-Q by KV heads, index-K replicated.""" + q, k, v, iq, ik = t + return torch.cat( + [_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0 + ).contiguous() + + +def _qkv_sparse_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v, iq, ik = t + return _swizzle( + torch.cat([_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0) + ) + + +def _qkv_dense_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +def _qkv_dense_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return _swizzle(torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0)) + + +def _index_qk_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[index-Q | index-K]`` in BF16: this rank's index-Q heads and the replicated index-K, the E4M3 values cast + (the scale tensors, the last two sources, are consumed only).""" + iq, ik = t[0], t[1] + return torch.cat([_col(iq, d), _dev(ik, d)], dim=0).to(d.dtype).contiguous() + + +def _gate_up_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return torch.cat([_col(gate, d), _col(up, d)], dim=0).contiguous() + + +def _gate_up_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return _swizzle(torch.cat([_col(gate, d), _col(up, d)], dim=0)) + + +def _row_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _row(t, d).contiguous() + + +def _row_s(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _swizzle(_row(t, d)) + + +def _col_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _col(t, d).contiguous() + + +#: Permute indices of the trtllm-gen expert relayouts, keyed by (shape, family, scale group): one expert's are +#: every expert's, so each is computed once per load. +_PERMUTE_CACHE: dict = {} + + +def _expert_w31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 operand: ``[up ; gate]`` packed-FP4 rows of this rank, then the gated-activation row + reorder and the epilogue-tile shuffle.""" + up, gate = t + w = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + w, _PERMUTE_CACHE, EPILOGUE_TILE_M, is_gated_act_gemm=True + ) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 E4M3 scales: ``[up ; gate]`` rows of this rank, the same row reorder at scale + granularity, then the 128x4 swizzle.""" + up, gate = t + s = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK, is_gated_act_gemm=True + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled.view(torch.uint8).reshape(shape)) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + + +def _expert_w2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 operand: this rank's input columns, then the epilogue-tile row shuffle.""" + w = _row(t, d).contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(w, _PERMUTE_CACHE, EPILOGUE_TILE_M) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 E4M3 scales: this rank's columns, the row shuffle at scale granularity, the swizzle.""" + s = _row(t, d).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w2_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled).view(torch.float8_e4m3fn).reshape(shape) + ) + + +# ------------------------------------------------------------------------------------------------ sources + + +def _one(template: str): + return lambda d, i, t=template: [(_LM + t.format(i=i), None)] + + +def _many(*templates: str): + return lambda d, i, ts=templates: [(tuple(_LM + t.format(i=i) for t in ts), None)] + + +def _per_expert(*suffixes: str): + """One source per routed expert, written at its expert index (all experts are local at tp4).""" + + def build(d, i, sx=suffixes): + out = [] + for e in range(d.num_experts): + keys = tuple(_LM + f"model.layers.{i}.block_sparse_moe.experts.{e}.{s}" for s in sx) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _dense_layers(d): + return range(d.first_sparse) + + +def _sparse_layers(d): + return range(d.first_sparse, d.num_layers) + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense layers, and every layer without the five-way packed + projection (``fuse_qkv_index_projection=False``).""" + return range(d.first_sparse) if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return range(d.first_sparse, d.num_layers) if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else range(d.first_sparse, d.num_layers) + + +_ATTN = "model.layers.{i}.self_attn" +_MOE = "model.layers.{i}.block_sparse_moe" + + +class MiniMaxM3Tp4Weights: + """MiniMax-M3 NVFP4 / sm_107 / tp4: the language model's table and the passes over it. + + ``declare`` allocates every weight it names, ``manifest`` maps checkpoint keys onto them, ``load`` fills + them; ``_targets`` is the one place that decides which parameters exist. + """ + + #: Free the per-expert relayout scratch after each expert stack. + RELEASE_AFTER = "_e_s2" + + WEIGHTS: tuple[W, ...] = ( + # --- the two boundary norms of every layer (Gemma, folded) --- + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + transform=_gemma_fold, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + transform=_gemma_fold, + ), + # --- attention: per-head Q/K norms (Gemma, unfolded), the packed QKV and the output projection --- + W("q_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".q_norm.weight")), + W("k_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".k_norm.weight")), + W( + "qkv_w", + shape=lambda d: (d.qkv_dense_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in "qkv")), + transform=_qkv_dense_w, + layers=_three_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_dense_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many(*(f"{_ATTN}.{p}_proj.weight_scale_inv" for p in "qkv")), + transform=_qkv_dense_s, + layers=_three_way_qkv_layers, + ), + W( + "qkv_w", + shape=lambda d: (d.qkv_sparse_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in ("q", "k", "v", "index_q", "index_k"))), + transform=_qkv_sparse_w, + layers=_five_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_sparse_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + *( + f"{_ATTN}.{p}_proj.weight_scale_inv" + for p in ("q", "k", "v", "index_q", "index_k") + ) + ), + transform=_qkv_sparse_s, + layers=_five_way_qkv_layers, + ), + # The separate BF16 index projection [index-Q | index-K] of the unfused sparse layers. The built-in model + # builds it as an unquantized Linear and fills it through the fused gate/up loader, which casts the + # checkpoint's E4M3 values to BF16 and does not apply their block scales: so does this table (the scale + # tensors are consumed, not applied). + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + src=_many( + _ATTN + ".index_q_proj.weight", + _ATTN + ".index_k_proj.weight", + _ATTN + ".index_q_proj.weight_scale_inv", + _ATTN + ".index_k_proj.weight_scale_inv", + ), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_layers, + ), + W( + "o_w", + shape=lambda d: (d.hidden, d.q_width), + dtype=torch.float8_e4m3fn, + src=_one(_ATTN + ".o_proj.weight"), + transform=_row_w, + ), + W( + "o_s", + shape=lambda d: (d.hidden * (d.q_width // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_ATTN + ".o_proj.weight_scale_inv"), + transform=_row_s, + ), + # --- dense SwiGLU-OAI MLP (layers 0..2) --- + W( + "mlp_gu_w", + shape=lambda d: (2 * d.dense_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight", "model.layers.{i}.mlp.up_proj.weight" + ), + transform=_gate_up_w, + layers=_dense_layers, + ), + W( + "mlp_gu_s", + shape=lambda d: (2 * d.dense_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight_scale_inv", + "model.layers.{i}.mlp.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_dense_layers, + ), + W( + "mlp_dn_w", + shape=lambda d: (d.hidden, d.dense_inter), + dtype=torch.float8_e4m3fn, + src=_one("model.layers.{i}.mlp.down_proj.weight"), + transform=_row_w, + layers=_dense_layers, + ), + W( + "mlp_dn_s", + shape=lambda d: (d.hidden * (d.dense_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one("model.layers.{i}.mlp.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_dense_layers, + ), + # --- MoE (layers 3..59): the fp32 router and its bias, the shared expert --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + dtype=torch.float32, + src=_one(_MOE + ".gate.weight"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "sh_gu_w", + shape=lambda d: (2 * d.shared_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + _MOE + ".shared_experts.gate_proj.weight", _MOE + ".shared_experts.up_proj.weight" + ), + transform=_gate_up_w, + layers=_sparse_layers, + ), + W( + "sh_gu_s", + shape=lambda d: (2 * d.shared_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _MOE + ".shared_experts.gate_proj.weight_scale_inv", + _MOE + ".shared_experts.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_sparse_layers, + ), + W( + "sh_dn_w", + shape=lambda d: (d.hidden, d.shared_inter), + dtype=torch.float8_e4m3fn, + src=_one(_MOE + ".shared_experts.down_proj.weight"), + transform=_row_w, + layers=_sparse_layers, + ), + W( + "sh_dn_s", + shape=lambda d: (d.hidden * (d.shared_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_MOE + ".shared_experts.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_sparse_layers, + ), + # --- the routed experts, NVFP4 (trtllm-gen layouts), all 128 on every rank --- + W( + "e_w31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_per_expert("w3.weight", "w1.weight"), + transform=_expert_w31, + layers=_sparse_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight"), + transform=_expert_w2, + layers=_sparse_layers, + ), + W( + "e_s31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w3.weight_scale", "w1.weight_scale"), + transform=_expert_s31, + layers=_sparse_layers, + ), + W( + "e_s2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w2.weight_scale"), + transform=_expert_s2, + layers=_sparse_layers, + ), + # The per-tensor scalars of every expert, raw; modeling derives the kernels' scalars from them. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + layers=_sparse_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_scale"), + ("isc3", "w3.input_scale"), + ("isc2", "w2.input_scale"), + ("ws2_1", "w1.weight_scale_2"), + ("ws2_3", "w3.weight_scale_2"), + ("ws2_2", "w2.weight_scale_2"), + ) + ), + # --- not per-layer --- + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=_one("model.embed_tokens.weight"), + layers=None, + ), + W( + "final_norm", + shape=lambda d: (d.hidden,), + src=_one("model.norm.weight"), + transform=_gemma_fold, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + sparse = cfg.sparse_attention_config + freq = [int(f) for f in sparse["sparse_attention_freq"]] + tp = mapping.tp_size + head_dim = int(cfg.head_dim) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + index_heads = max(int(sparse["sparse_num_index_heads"]) // tp, 1) + index_dim = int(sparse["sparse_index_dim"]) + return SimpleNamespace( + fused=bool( + getattr( + core.model_config.sparse_attention_config, "fuse_qkv_index_projection", False + ) + ), + num_layers=cfg.num_hidden_layers, + first_sparse=freq.index(1), + hidden=cfg.hidden_size, + head_dim=head_dim, + heads=heads, + kv_heads=kv_heads, + index_heads=index_heads, + index_dim=index_dim, + q_width=heads * head_dim, + qkv_dense_rows=(heads + 2 * kv_heads) * head_dim, + qkv_sparse_rows=(heads + 2 * kv_heads) * head_dim + index_heads * index_dim + index_dim, + dense_inter=int(cfg.dense_intermediate_size) // tp, + shared_inter=int(cfg.shared_intermediate_size) * int(cfg.n_shared_experts) // tp, + moe_inter=int(cfg.intermediate_size) // tp, + num_experts=int(cfg.num_local_experts), + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None)] + return [(f"l{i}_{entry.name}", i) for i in layers] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares (meta under the engine's meta init).""" + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key or key tuple, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, core, weights) -> set: + """The vision tower and its projectors, and the multi-token-prediction head: in the checkpoint, never built + by this target.""" + return {k for k in weights.keys() if k.startswith(_UNBUILT_PREFIXES)} + + def load(self, model, weights) -> None: + """Fill every declared weight and the shell's ``lm_head`` shard; assert coverage both ways.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + assert set(manifest) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest) ^ set(core.w.keys()), + ) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(core, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(core.w[param_key], ckpt, index, transform, weights, d, consumed) + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + # The shell owns the vocab-parallel head: this rank's rows of the full table. + _fill(model.lm_head.weight, _LM + "lm_head.weight", None, _col_w, weights, d, consumed) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + } + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the unbuilt families (vision, MTP head); " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _materialize(entry) -> torch.Tensor: + """``[:]`` on a lazy safetensors slice; a 0-dim tensor (the per-tensor scalars) rejects that index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +def _fill(param, ckpt, index, transform, weights, d, consumed) -> None: + keys = ckpt if isinstance(ckpt, tuple) else (ckpt,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + d.dev = dst.device + # Host tensors: a transform slices this rank's part on the host and moves only that to the device. + src_all = tuple(_materialize(weights[key]) for key in keys) + if transform is not None: + src = transform(src_all if isinstance(ckpt, tuple) else src_all[0], d) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + if src.dim() == 0 and dst.dim() == 0: + pass + elif src.numel() == 1 and dst.numel() == 1: + src = src.reshape(dst.shape) + src = src.to(dst.device, non_blocking=True) + assert tuple(dst.shape) == tuple(src.shape), (ckpt, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + _release(weights, keys) + + +def _release(weights, keys) -> None: + """Drop consumed checkpoint tensors from the engine's dict as the load goes (bounded host memory).""" + if hasattr(weights, "mark_consumed_keys"): + weights.mark_consumed_keys(keys) + else: + for key in keys: + weights.pop(key, None) + + +MODEL_WEIGHTS = MiniMaxM3Tp4Weights() + + +# -------------------------------------------------------------------------------------------- the drafter + + +def _draft_one(template: str): + return lambda d, i, t=template: [(t, None)] + + +def _draft_many(*templates: str): + return lambda d, i, ts=templates: [(tuple(ts), None)] + + +def _draft_fc_norms(d): + return range(d.capture_layers) + + +def _draft_qkv(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +class MiniMaxM3Tp4DraftWeights: + """The one-layer BF16 Eagle3 drafter (``speculative_model``), tp4. + + The fc over the three captured target layers and its per-layer fc_norms, the layer's three norms, its QKV + (input ``[embeds | hidden]``, so twice the hidden width; Q by heads, K/V by KV heads) and o_proj, its gated + MLP, the final norm and the full-vocabulary head (vocab-parallel). The checkpoint's ``embed_tokens`` is the + target's table (the drafter reads the target's); it is a predicted non-load. + """ + + WEIGHTS: tuple[W, ...] = ( + W( + "fc", + shape=lambda d: (d.hidden, d.capture_layers * d.hidden), + src=_draft_one("fc.weight"), + layers=None, + ), + W( + "fc_norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"fc_norm.{i}.weight", None)], + layers=_draft_fc_norms, + ), + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.input_layernorm.weight"), + layers=None, + ), + W( + "hidden_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.hidden_norm.weight"), + layers=None, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.post_attention_layernorm.weight"), + layers=None, + ), + W( + "qkv", + shape=lambda d: (d.qkv_rows, 2 * d.hidden), + src=_draft_many(*(f"layers.0.self_attn.{p}_proj.weight" for p in "qkv")), + transform=_draft_qkv, + layers=None, + ), + W( + "o", + shape=lambda d: (d.hidden, d.q_width), + src=_draft_one("layers.0.self_attn.o_proj.weight"), + transform=_row_w, + layers=None, + ), + W( + "gu", + shape=lambda d: (2 * d.inter, d.hidden), + src=_draft_many("layers.0.mlp.gate_proj.weight", "layers.0.mlp.up_proj.weight"), + transform=_gate_up_w, + layers=None, + ), + W( + "dn", + shape=lambda d: (d.hidden, d.inter), + src=_draft_one("layers.0.mlp.down_proj.weight"), + transform=_row_w, + layers=None, + ), + W("final_norm", shape=lambda d: (d.hidden,), src=_draft_one("norm.weight"), layers=None), + ) + + def dims(self, draft) -> SimpleNamespace: + cfg = draft.draft_config.pretrained_config + mapping = draft.draft_config.mapping + tp = mapping.tp_size + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + return SimpleNamespace( + hidden=cfg.hidden_size, + capture_layers=draft.capture_layers, + head_dim=head_dim, + q_width=heads * head_dim, + qkv_rows=(heads + 2 * kv_heads) * head_dim, + inter=cfg.intermediate_size // tp, + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + if entry.layers is None: + return [(entry.name, None)] + return [(f"{entry.name}{i}", i) for i in entry.layers(d)] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, draft) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + d = self.dims(draft) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, draft, weights) -> set: + """The drafter checkpoint's own embedding: the drafter reads the target's table.""" + return {"embed_tokens.weight"} & set(weights.keys()) + + def load(self, draft, weights) -> None: + d = self.dims(draft) + manifest = self.manifest(draft) + assert set(manifest) == set(draft.w.keys()), set(manifest) ^ set(draft.w.keys()) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(draft, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(draft.w[param_key], ckpt, index, transform, weights, d, consumed) + _fill(draft.lm_head.weight, "lm_head.weight", None, _col_w, weights, d, consumed) + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + f"drafter checkpoint coverage: unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = MiniMaxM3Tp4DraftWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py new file mode 100644 index 000000000000..adf71d21a10a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107 / tp8.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/__init__.py new file mode 100644 index 000000000000..c5fd79e59d7f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107 / tp8 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp8/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp8/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp8/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/long_select.py new file mode 100644 index 000000000000..14e24fbdbcac --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/long_select.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. + +Same contract and bit-identical output as ``torch.ops.trtllm.minimax_m3_select_blocks`` on per-KV-head scores +[num_kv_heads, n_blocks, total_q]: init / local forcing, per-query valid-block count, ties to the smaller block index, +forced -inf picks emitted as -1, ascending block ids with -1 tail padding. + +One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups ``block % 16`` is a lower +bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared memory +and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every candidate +against the whole row (re-read). + +The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long`` (this target's copy of ``minimax_m3_select_long`` +with a choice of where it releases the sparse decode), called through this target's ``m3_select_long`` entry, or with +``dsl`` the same selection from one 384-thread CTA per row (``m3_select_long_dsl``, kernels/m3_select_long_dsl). +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long +from ..ops.attention.m3_select_long_dsl import m3_select_long_dsl + + +def long_select_blocks( + max_score_kv: torch.Tensor, + *, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + dsl: bool = False, +) -> Optional[torch.Tensor]: + """Top-16 select over per-KV-head scores [num_kv_heads, n_blocks, total_q] with 129..9216 blocks. + + Returns the head-major backed [total_q, num_kv_heads, 16] int32 table, or None outside this kernel's shapes + (callers keep the reference op). ``early_trigger`` releases the dependent sparse decode right after the select's + grid wait (modeling's ``SELECT_EARLY_MAX_ROWS``); otherwise ``trigger_phase`` 1 releases it once every score of the + row is loaded and 0 at the select's exit (modeling's ``SELECT_LONG_TRIGGER_PHASE``). ``dsl`` runs the 384-thread + kernel (modeling's ``SELECT_LONG_DSL``; the same table). + """ + if ( + max_score_kv.dtype != torch.float32 + or not max_score_kv.is_cuda + or max_score_kv.dim() != 3 + or not 128 < max_score_kv.shape[1] <= MAX_BLOCKS + or n_valid_blocks.dtype != torch.int32 + or not n_valid_blocks.is_cuda + or n_valid_blocks.shape[0] < max_score_kv.shape[2] + ): + return None + num_kv_heads, _, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, 16), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + return (m3_select_long_dsl if dsl else m3_select_long)( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_attn_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_attn_lc8.py new file mode 100644 index 000000000000..2d446f2b45d3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_attn_lc8.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 long-context sparse decode attention at 8 query heads in one launch. + +Replaces the generic chain's split-K producer (``_gqa_sparse_decode_kernel``, one block per chunk) and the LSE +merge (``_merge_topk_attn_out_kernel``, with its MXFP8 copy for the o_proj) once the top-k block table exists. + +Grid: one cluster of 16 CTAs per query token; CTA ``r`` attends the token's ``r``-th selected block for all 8 heads, +keeps the partial in shared memory, and the cluster merges the 16 partials over DSMEM into the BF16 row and, +optionally, its MXFP8 copy and swizzled UE8M0 scales exactly as ``_store_mxfp8_head`` lays them out. + +Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and PV +on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +log-domain recurrence. Accumulation order inside the MMAs and the reductions differs, so results match to rounding, +not bit for bit. + +Schedule: the kernel releases its PDL dependents right after its grid wait, so the o_proj launches and streams its +weights while attention runs. Up to ``PREWAIT_MAX_ROWS`` query rows it also stages the request length, +the query rows and the request's block-table row before that wait, overlapping them with the select ahead when it +releases this grid early (modeling's ``SELECT_EARLY_MAX_ROWS``). Neither changes a value. + +Geometry: 8 query heads per KV head (TP8), one KV head, head_dim 128, page 128, top-k 16. The kernel is +``torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8``, called through this target's ``m3_sparse_decode_lc8`` entry. +""" + +import torch + +from ..ops.attention.m3_sparse_decode_lc8 import NUM_HEADS, WIDTH, m3_sparse_decode_lc8 + +# Up to this many query rows the kernel stages the request length, the query rows and the block-table row before its +# grid wait. +PREWAIT_MAX_ROWS = 8 # capacity: fixed: measured occupancy crossover +# The split-K merge forms its weights once per CTA (MERGE_ONCE: also enables the identity route and the L2 prefetches) +# and divides once per output (ONE_DIV_MERGE: sum_c w_c acc_c / sum_c w_c l_c in chunk order; deterministic, in the +# float64 reference's band). +MERGE_ONCE = True +ONE_DIV_MERGE = True + + +class PrefetchSlot: + """The read-only tensors the next long-context sparse decode prefetches into L2 (the attention layer's o_proj weight + and scales), set by the layer around its decode call and taken by the decode; empty means no prefetch.""" + + def __init__(self) -> None: + self.tensors = None + + def take(self): + tensors, self.tensors = self.tensors, None + return tensors + + +PREFETCH = PrefetchSlot() +# (this layer's id buffer, the next sparse layer's id buffer, its K / V pools, its K / V block scales) +NEXT_KV = PrefetchSlot() + + +def lc8_applies(q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output) -> bool: + """Whether the step's operands fit the kernel: TP8 geometry, packed NVFP4 pages, FP8 q, BF16 [T, 8, 128] out.""" + if q.dim() != 3 or tuple(q.shape[1:]) != (NUM_HEADS, 128) or q.dtype != torch.float8_e4m3fn: + return False + if q.stride(2) != 1 or q.stride(1) % 16 or q.stride(0) % 16 or q.data_ptr() % 16: + return False + if int(k_paged.shape[1]) != 1 or int(topk_idx.shape[0]) != 1 or int(topk_idx.shape[-1]) != 16: + return False + for t in (k_paged, v_paged): + if ( + t.stride(3) != 1 + or t.stride(2) != 64 + or int(t.shape[3]) != 64 + or t.data_ptr() % 16 + or t.stride(0) % 16 + ): + return False + for t in (k_block_scale, v_block_scale): + if t is None or t.stride(3) != 1 or t.stride(2) != 8 or t.data_ptr() % 8 or t.stride(0) % 8: + return False + return ( + output.dtype == torch.bfloat16 + and output.dim() == 3 + and output.is_contiguous() + and tuple(output.shape) == tuple(q.shape) + ) + + +def m3_lc8_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + linear_softmax=True, + mxfp8=None, + use_pdl=True, + identity=False, + identity_step=False, + prefetch=None, + next_kv=None, +): + """Drop-in for ``minimax_m3_sparse_attn_decode`` on the TP8 NVFP4 shape; callers check ``lc8_applies`` first. + + ``topk_idx`` is the head-major [1, total_q, 16] view the Triton producer takes. ``mxfp8`` is ``(q8, sf)`` from + ``MERGE_MXFP8.buffers(1024, dev)`` to also emit the o_proj hand-off. With ``MERGE_ONCE``, + ``prefetch`` (up to two contiguous tensors) is pulled into L2 by the op's CTAs before their grid dependency, and + ``next_kv`` drives the next-layer KV prefetch. + """ + merge_once = MERGE_ONCE + total_q = int(q.shape[0]) + q8, sf = ( + (mxfp8[0].view(torch.uint8), mxfp8[1].view(torch.uint8)) + if mxfp8 is not None + else (None, None) + ) + m3_sparse_decode_lc8( + q.view(torch.uint8), + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + output.view(total_q, WIDTH), + q8, + sf, + WIDTH // 128, + int(decode_query_len), + float(sm_scale), + bool(linear_softmax), + bool(use_pdl), + total_q <= PREWAIT_MAX_ROWS, + merge_once=merge_once, + identity=bool(identity) and merge_once, + identity_step=bool(identity_step) and merge_once, + one_div=ONE_DIV_MERGE and merge_once, + prefetch=prefetch if merge_once else None, + next_kv=( + tuple(next_kv[:2]) + tuple(t.view(torch.uint8) for t in next_kv[2:]) + if merge_once and next_kv is not None + else None + ), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_draft_step.py new file mode 100644 index 000000000000..c977362e92b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_draft_step.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 draft pass for the MiniMax-M3 drafter's vocab-sharded BF16 head (TP decode, <= 8 requests). + +One kernel replaces a draft pass's head and the draft loop's per-pass bookkeeping, at the in-tree precision: + +* The drafter's BF16 lm_head shard is kept a second time, built by ``prepare_head`` before the executor's warmup: in + the lossless 12-bit coding of the target head (``ops.gemm.m3_target_head``; 75% of the bytes, kept only when it + decodes to the BF16 shard bit for bit), else in ``mma.m16n8k16`` A-fragment order (``bf16_fragments``, the same + values). The kernel streams it with BF16 MMAs and FP32 accumulation over the BF16 head input rows; a coded + exception decodes to +0 and its exact product x * w is added at its tile's flush. +* Each logit is rounded to BF16 (the in-tree logits are the BF16 head GEMM's output) and keyed + ``ordered(value) << 32 | ~global_index`` with -0 as +0, so the larger value wins and ties take the lowest index, as + the in-tree ``torch.max`` / ``torch.argmax`` over the gathered logits (``greedy_sample_draft_with_tp_gather``). +* The last CTA exchanges each row's key across the TP ranks through a multicast mapping (Lamport slots, see + ``m3_shard_accept``), then stores the draft token, the next pass's input id and position, and advances the draft KV + lengths as the in-tree draft loop does between passes. + +The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu) and called through the +catalog entry ``ops.spec.m3_draft_step``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok +from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments, max_tile_pieces +from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op +from . import m3_report +from .m3_shard_accept import AcceptWorkspace, create_workspace + +_ENGAGED: set = set() + + +# The fused draft loop streams the 12-bit coded copy of the draft head when it decodes bit for bit; off (the coded +# head is slower on sm_107: the BF16 fragments stream). +HEAD_CODED = False +# The head copy placed per locality domain, each pass as two launches: off at TP8 (one launch over the whole device). +UGPU = False +# The draft step's CTA count (at most one per SM; 0 = one per SM). +GRID = 208 # capacity: fixed: a launch-grid choice, not a size +# The draft head sums a tile shared by more than 2 warp pieces in warp order instead of with atomics (the same sum +# every run; the atomics are deterministic only where a tile takes at most 2 pieces). +FIXED_ORDER = True +FIXED_ORDER_PATH = "m3_draft_step_fixed_order" + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the 12-bit coded stream + with its exception lists (``coded``), or the fragment-ordered BF16 copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + coded: bool = False + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + exp_base: int = 0 + # BF16 fragments placed per locality domain (UGPU below): tiles [0, split_tile) live in domain 0 and + # the rest in domain 1; 0 = not placed (one launch over all SMs). + split_tile: int = 0 + + +def head_supported(lm_head) -> bool: + """A bias-free BF16 head whose local shard the kernel can own: whole 16-row tiles, K a multiple of 16, at most + MAX_TILES_PER_CTA tiles per CTA of a one-CTA-per-SM grid.""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + if getattr(lm_head, "bias", None) is not None: + return False + n, k = weight.shape + if n % 16 or k % 16: + return False + return (n // 16 + _grid() - 1) // _grid() <= MAX_TILES_PER_CTA + + +def prepare_head(lm_head, tp_rank: int) -> DraftHead: + """The streamed copy of ``lm_head``'s local shard (global ids from tp_rank * rows): its 12-bit coding when + ``HEAD_CODED``, K is a multiple of 64 and the coding decodes to the shard bit for bit, else (fail + closed) the fragment-ordered BF16 copy. + Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_draft_head``).""" + weight = lm_head.weight + offset = tp_rank * weight.shape[0] + with torch.no_grad(): + w = weight.detach() + if HEAD_CODED and w.shape[1] % 64 == 0: + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return DraftHead( + weight=weight, + wf=packed, + vocab_offset=offset, + coded=True, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + exp_base=base, + ) + del packed, exc_off, exc_pos, exc_w + wf = bf16_fragments(w) + wf, split_tile = _place_per_domain(wf) + return DraftHead(weight=weight, wf=wf, vocab_offset=offset, split_tile=split_tile) + + +def _place_per_domain(wf: torch.Tensor): + """(wf, split tile): the BF16 fragments split-mapped over the two locality domains at the granule-aligned tile + nearest half (locality.LocalityMap.place_split; its domain streams built here, before capture), or (wf, 0) when + UGPU is off. Placement only: the streamed values are the same.""" + if not UGPU: + return wf, 0 + from .m3_locality import LocalityMap + + lm = LocalityMap.get(wf.device) + tile_bytes = wf[0].numel() * wf.element_size() + step = 1 + while (step * tile_bytes) % lm.granularity: + step += 1 + n_tiles = wf.shape[0] + split_tile = max(step, round(n_tiles / 2 / step) * step) + per = lm.sms_per_domain + if split_tile >= n_tiles or max(split_tile, n_tiles - split_tile) > per * MAX_TILES_PER_CTA: + return wf, 0 + lm.domain_streams() + return lm.place_split(wf, split_tile), split_tile + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _grid() -> int: + """The draft step's CTA count: GRID when set (at most one per SM), else one per SM.""" + return min(GRID, _num_sms()) if GRID > 0 else _num_sms() + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace ([2][MAX_TOKENS][world] slots), created collectively on the first eager call; a + CUDA-graph capture that finds it missing is a miss.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("m3_draft_step", None, "the draft-step exchange workspace") + return None + _WS = create_workspace(mapping, device, MAX_TOKENS) + return _WS + + +def _fixed_order( + rows: int, head: DraftHead, x: torch.Tensor, num_ctas: int, split: bool = False +) -> bool: + """FIXED_ORDER where the launch lets a tile take more than 2 warp pieces (max_tile_pieces; a + split or equal-chunk launch always), with its engagement line (D-FIX); False keeps the atomics, deterministic on + this shape.""" + if not FIXED_ORDER: + return False + eq = getattr(head, "bnd", None) is not None + if not (split or eq): + chunks = x.shape[1] // (16 if head.exc_off is None else 64) + if ( + max_tile_pieces( + head.wf.shape[0] if head.exc_off is None else head.exc_off.numel() - 1, + num_ctas, + chunks, + ) + <= 2 + ): + return False + m3_report.taken(FIXED_ORDER_PATH, rows, "the draft head sums shared tiles in warp order") + return True + + +def draft_step( + head: DraftHead, + x: torch.Tensor, + rows: int, + x_rows: Optional[torch.Tensor], + mapping, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + fold: Optional[tuple] = None, +) -> bool: + """One fused greedy draft pass over ``rows`` requests (see the module docstring); False when the exchange + workspace is missing during a capture. Collective over the TP group. + + Args: + head: the prepared head of this rank (``prepare_head``), coded or BF16 fragments. + x: [*, K] BF16 head input rows (the drafter's final-norm output); row ``x_rows[b]`` (int64) or ``b`` is + request b's. + draft_out: int32 [rows, >= step + 1], unit column stride: column ``step`` receives the tokens. + ids_next / pos_next: int32 [rows]: the tokens, and ``pos_src[x_rows[b] or b] + 1`` (pos_src int32). + kv_lens: int32 draft KV lengths: ``-= draft_len - num_acc`` with ``seq_lens = 1`` on pass 0 when ``num_acc`` + is given, ``+= 1`` on later passes. + fold: on the last pass (modeling's DRAFT_STEP_FOLD_TAIL), (the 14 fold-tail tensors, rows, tokens per step) from + M3SpecStepFold.fused_tail_operands: the step's forward-tail store and next-step staging run by this pass's + last CTA after its tokens, instead of the standalone fold tail. + """ + ws = workspace(mapping, x.device) + if ws is None: + return False + if head.split_tile: + _draft_step_split( + head, + x, + rows, + x_rows, + ws, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + fold, + ) + return True + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + _grid(), + True, + None, + head.exc_off, + head.exc_pos, + head.exc_w, + head.exp_base, + fold=fold, + fixed_order=_fixed_order(rows, head, x, _grid()), + ) + return True + + +def _draft_step_split( + head, + x, + rows, + x_rows, + ws, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + fold=None, +) -> None: + """One pass as two launches, one per locality domain on that domain's green-context stream, each streaming its own + domain's tiles of the split-mapped fragments; the keys' max and the last-CTA exchange and bookkeeping run once over + both (the last of the 2 x SMs-per-domain CTAs). Forked from and joined back to the current stream; no PDL on the + forks.""" + from .m3_locality import LocalityMap + + lm = LocalityMap.get(x.device) + per = lm.sms_per_domain + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + streams = lm.domain_streams() + n_tiles = head.wf.shape[0] + if x.device not in _ENGAGED: + _ENGAGED.add(x.device) + logger.info( + f"[M3 sm_107 tp2] fast path engaged: m3_draft_step_ugpu (split tile {head.split_tile} of {n_tiles}, " + f"{per} CTAs per domain, on {x.device})" + ) + for st, lo, hi in ((streams[0], 0, head.split_tile), (streams[1], head.split_tile, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + per, + False, + None, + None, + None, + None, + 0, + lo, + hi, + 2 * per, + fold=fold, # run by whichever launch holds the pass's last CTA + fixed_order=_fixed_order(rows, head, x, 2 * per, split=True), + ) + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_drafter_attn.py new file mode 100644 index 000000000000..782ff2afb816 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_drafter_attn.py @@ -0,0 +1,404 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's decode attention on Rubin: ``M3DrafterFullCtxFmha``. + +The drafter layer is a plain ``TrtllmAttention`` over its own FP8 paged cache (the full context, every draft pass). +On SM 10.7 its generation phase runs ``thop.attention``'s in-tree trtllm-gen runner, whose split is capped at 8 CTAs +per (request, KV head): 45-46 us per pass at 165k tokens on 212 SMs, for 85 MB. + +This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append +into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context +PrimTS split-KV decode of ``kernels/m3_fullctx_decode`` at the same pinned recipes as the dense layers. + +With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead +(``minimax_m3_sm107tp8_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page +pool and block table taken from the KV cache pool pointers and mapping the preprocess reads. + +It serves generation-only steps of the covered (query rows per request, requests) shapes; the target installs it +first in the drafter layer's FMHA list (``install``), so every other step falls through to the stock libraries. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + FlashInferTrtllmGenFmha, +) +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode + +from . import m3_fullctx_decode as fullctx + +# With PDL, the fused RoPE + FP8 K/V launch takes the maximum shared-memory carveout so the drafter decode's large-smem +# CTAs co-reside with it (timing only). +ROPE_CARVEOUT = True + + +def fwd_has_no_mrope(fwd) -> bool: + return getattr(fwd, "mrope_position_deltas", None) is None + + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + +class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): + """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + + def __init__( + self, + attn: "TrtllmAttention", + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + pdl: bool = False, + on_rope_taken=None, + ) -> None: + super().__init__(attn) + self._recipes = recipes + self._max_kv_len = max_kv_len + self._own_reducer = own_reducer + self._fused_rope_mode = fused_rope_mode + self._pdl = pdl + self._on_rope_taken = on_rope_taken + self._on_taken = on_taken + self._on_uncovered = on_uncovered + self._on_missed = on_missed + + def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: + nreq = int(metadata.num_generations) + if nreq <= 0 or int(q.shape[0]) % nreq: + return None + return int(q.shape[0]) // nreq, nreq + + def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: + if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: + return False + if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: + return False + attn = self.attn + if ( + attn.num_heads != fullctx.HEADS + or attn.num_kv_heads != fullctx.KV_HEADS + or attn.head_dim != fullctx.HEAD_DIM + or not QuantMode(attn.quant_mode).has_fp8_kv_cache() + ): + self._on_uncovered("not the TP8 FP8 drafter geometry") + return False + # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's + # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. + if ( + forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant + or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig + ): + self._on_uncovered("per-step KV cache scales") + return False + if int(metadata.max_seq_len) > self._max_kv_len: + self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") + return False + shape = self._shape(q, metadata) + if shape is None or shape not in self._recipes: + self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") + return False + # Compiled before warmup (``install``); the forward never compiles one. + if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: + self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") + return False + return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) + + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + attn = params.attn + width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim + rotary_dim = int(attn.rope_params.dim) + return ( + self._fused_rope_mode is not None + and params.qkv_input is not None + and attention_input.dim() == 2 + and attention_input.dtype == torch.bfloat16 + and attention_input.is_contiguous() + and int(attention_input.shape[1]) == width + and not params.is_cross + and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) + and attn.rotary_cos_sin is not None + and attn.rotary_cos_sin.dtype == torch.float32 + and fwd_has_no_mrope(params.fwd) + and 0 < rotary_dim <= attn.head_dim + and rotary_dim % 16 == 0 + and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) + # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo + # runs for several requests only) and reads the workspace's, so it keeps that path. + and not (self._multi_token(params) and int(params.batch_size) == 1) + ) + + @staticmethod + def _multi_token(params: FmhaParams) -> bool: + return ( + params.spec_decoding_generation_lengths is not None + and int(params.attn.predicted_tokens_per_seq) > 1 + ) + + def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" + if not self._fused_rope_applies(params, attention_input): + return False + attn = params.attn + meta = params.meta + qlen, nreq = int(params.input_seq_length), int(params.batch_size) + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + if plan is None: + return False + # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary + # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the + # block table is the pool's [requests, 2, blocks] K / V sub-page indices. + pool_index, layer_in_pool = ( + int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] + ) + tpb = int(params.tokens_per_block) + sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim + base = ( + int(meta.host_kv_cache_pool_pointers[pool_index, 0]) + + layer_in_pool * int(params.kv_factor) * sub_page_bytes + ) + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) + ) + tables = meta.kv_cache_block_offsets[ + pool_index, params.seq_offset : params.seq_offset + nreq + ] + multi = self._multi_token(params) + if not torch.cuda.is_current_stream_capturing(): + # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess + # path checks; the op itself writes V at the table's V sub-page. + tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(tables.shape[-1], device=tables.device) + in_use = cols[None, :] < tpb_held[:, None] + offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + q = torch.empty( + int(params.num_tokens), + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=attention_input.device, + ) + torch.ops.trtllm.minimax_m3_sm107tp8_drafter_rope_fp8_kv( + attention_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + int(attn.rope_params.dim), + tpb, + self._fused_rope_mode, + self._pdl, + ROPE_CARVEOUT, + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + tables[:, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + if self._on_rope_taken is not None: + self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") + self._on_taken(params.num_tokens) + return True + + def run_generation(self, params: FmhaParams) -> None: + attention_input = params.qkv_input if params.qkv_input is not None else params.query_input + if attention_input is None: + raise RuntimeError("the drafter decode requires QKV or query input.") + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + if self._run_fused_rope(params, attention_input): + return + # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 + # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). + fp8_context_fmha = False + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + _fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + _window_left, + is_multi_token_gen, + ) = thop.trtllm_gen_generation_preprocess( + attention_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + fp8_context_fmha, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + ) + qlen, nreq = params.input_seq_length, params.batch_size + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the + # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row + # pass is never multi-token. + if plan is None or (is_multi_token_gen and qlen == 1): + raise RuntimeError( + f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" + ) + rows = params.num_tokens * attn.num_heads * attn.head_dim + q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) + q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) + # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 + # of the pool (the two views have the same page count, as the kernel requires). + if not torch.cuda.is_current_stream_capturing(): + # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. + lengths = params.spec_decoding_generation_lengths + if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): + raise RuntimeError( + f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" + ) + # Only the blocks a request holds are laid out; the table's tail entries are unset. + tpb = int(params.tokens_per_block) + held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(block_tables.shape[-1], device=block_tables.device) + in_use = cols[None, :] < held[:, None] + offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. + # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + block_tables[:nreq, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + self._on_taken(params.num_tokens) + + +def install( + draft_model: torch.nn.Module, + recipes, + max_kv_len: int, + own_reducer: bool, + pdl: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + on_rope_taken=None, +) -> int: + """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention + backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it + was installed on.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + count = 0 + device = next(draft_model.parameters()).device + # The backend is not a module: each attention module of the drafter holds it as ``attn``. + for module in draft_model.modules(): + backend = getattr(module, "attn", None) + if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): + continue + if not bool((backend.kv_cache_scaling_factor == 1.0).all()): + on_uncovered("a scaled FP8 KV cache") + continue + fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) + libs = backend._fmha_manager.fmha_libs + if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): + libs.insert( + 0, + M3DrafterFullCtxFmha( + backend, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode=fused_rope_mode, + pdl=pdl, + on_rope_taken=on_rope_taken, + ), + ) + count += 1 + return count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fc_wait.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fc_wait.py new file mode 100644 index 000000000000..a545b43857ec --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fc_wait.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter fc fork's release of the first draft input by a device counter instead of a stream join (modeling's +``DI_FC_FLAG``). + +Triton. ``fc_wait_kernel``: one thread, queued on the main stream after the acceptance. It lets the next kernel (the +first draft input) launch, waits for the acceptance before it, then spins until the fork's counter reaches the number +of forks waited for. Ordering only: no value changes. +""" + +import torch +import triton +import triton.language as tl + +# Spins of fc_wait_kernel before it gives up on the fc fork (~ms): a broken fork / wait pairing then counts an overrun +# in counts[2] instead of hanging the step (the draft input may read a stale fc output: draft tokens only, the verify +# stays exact). +MAX_SPINS = 1 << 22 # capacity: fixed: a spin bound, not a size + + +@triton.jit +def fc_wait_kernel(counts, MAX_SPINS: tl.constexpr): + # counts[0]: fc forks finished (bumped on the aux stream after each fc); counts[1]: forks waited for (here). + # Lets the next kernel (the first draft input) launch now; waits for the acceptance before it, then for the fork. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + target = tl.load(counts + 1) + 1 + done = tl.load(counts, volatile=True) + spins = 0 + while (done < target) & (spins < MAX_SPINS): + done = tl.load(counts, volatile=True) + spins += 1 + tl.store(counts + 1, target) + if done < target: + tl.atomic_add(counts + 2, 1) + + +def fc_wait(counts: torch.Tensor) -> None: + """Queue the wait on the current stream for the fork ``counts`` (int64 [3]) last counted.""" + fc_wait_kernel[(1,)](counts, MAX_SPINS=MAX_SPINS, num_warps=1, launch_pdl=True) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fold_tail.py new file mode 100644 index 000000000000..292d1df264b1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fold_tail.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The speculative step's store and next-step staging at the tail of the forward (one Triton launch). + +For each of the first ``n`` rows (``n`` read from device memory, so a captured graph whose step was not armed writes +nothing), with ``s = slots[i]``: + +* the spec sampler's per-slot store, exactly as ``m3_spec_tail.slot_scatter`` writes it (zero-padded to each store + width): ``new[s][c]``, ``next_new[s][c]``, ``draft[s][c]`` and ``lens[s]``; +* the engine's next-step inputs, exactly as ``m3_spec_tail.step_gather`` writes them from that store for a step whose + rows are these rows in this order (no new tokens, first previous row 0): ``ids[i*T + t] = next_new[s][t]``, + ``draft_ids[i*D + d] = draft[s][d]``, ``pos_off[i*T + t] = lens[s]``, ``kv_off[i] = lens[s] - T``. + +Integer copies only: the values are the ones the sampler's and the engine's own copies write. + +With ``cpp`` the same writes run as this target's C++ op (trtllm::minimax_m3_sm107_tp8_fold_tail, +minimaxM3Sm107Tp8/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fold_tail_kernel( + ctrl, + slots, + src_new, + sn_s0, + sn_s1, + SN_W: tl.constexpr, + dst_new, + dn_s0, + dn_s1, + DN_W: tl.constexpr, + src_next, + sx_s0, + sx_s1, + SX_W: tl.constexpr, + dst_next, + dx_s0, + dx_s1, + DX_W: tl.constexpr, + src_draft, + sd_s0, + sd_s1, + SD_W: tl.constexpr, + dst_draft, + dd_s0, + dd_s1, + DD_W: tl.constexpr, + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + T: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(ctrl) + if i < n: + s = tl.load(slots + i).to(tl.int64) + c = tl.arange(0, BLOCK) + new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) + tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) + nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) + tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) + drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) + tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) + ln = tl.load(src_lens + i) + tl.store(dst_lens + s, ln) + # The next step's inputs read the store's padded rows: zero past each source width. + tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) + tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) + tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) + tl.store(kv_off + i, ln - T) + + +def fold_tail( + ctrl: torch.Tensor, + slots: torch.Tensor, + rows: int, + src_new: torch.Tensor, + dst_new: torch.Tensor, + src_next: torch.Tensor, + dst_next: torch.Tensor, + src_draft: torch.Tensor, + dst_draft: torch.Tensor, + src_lens: torch.Tensor, + dst_lens: torch.Tensor, + ids: torch.Tensor, + draft_ids: torch.Tensor, + pos_off: torch.Tensor, + kv_off: torch.Tensor, + tokens_per_step: int, + cpp: bool = False, + one_warp: bool = False, +) -> bool: + """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs + are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width]. + ``one_warp`` (with ``cpp``): the in-kernel form (one warp looping over the rows, foldTailWarp in + minimaxM3Sm107Tp8/foldTail.h) that a caller folding the tail into its own kernel runs; the same writes.""" + two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) + one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) + if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): + return False + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) + or slots.dtype != torch.long + ): + return False + t_step = tokens_per_step + d_step = t_step - 1 + if ( + rows <= 0 + or slots.numel() < rows + or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) + or ids.numel() < rows * t_step + or pos_off.numel() < rows * t_step + or draft_ids.numel() < rows * d_step + or kv_off.numel() < rows + ): + return False + widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) + widths += (src_draft.shape[1], dst_draft.shape[1], t_step) + if cpp and max(widths) <= 32: + op = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_fold_tail_one_warp + if one_warp + else torch.ops.trtllm.minimax_m3_sm107_tp8_fold_tail + ) + op( + ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, + src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, + ) # fmt: skip + return True + block = triton.next_power_of_2(max(widths)) + _fold_tail_kernel[(rows,)]( + ctrl, + slots, + src_new, + src_new.stride(0), + src_new.stride(1), + src_new.shape[1], + dst_new, + dst_new.stride(0), + dst_new.stride(1), + dst_new.shape[1], + src_next, + src_next.stride(0), + src_next.stride(1), + src_next.shape[1], + dst_next, + dst_next.stride(0), + dst_next.stride(1), + dst_next.shape[1], + src_draft, + src_draft.stride(0), + src_draft.stride(1), + src_draft.shape[1], + dst_draft, + dst_draft.stride(0), + dst_draft.stride(1), + dst_draft.shape[1], + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + t_step, + d_step, + block, + num_warps=1, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_force_vec.py new file mode 100644 index 000000000000..cc0ce28d9639 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_force_vec.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's +``force_num_accepted_tokens``), in one launch. Reached through the ``ops/spec`` entries ``m3_force_vec`` and +``m3_force_vec_req``. + +Triton. ``force_vec_kernel``: one CTA bumps the shared step counter and draws each request's count from the fixed-seed +pool with the in-tree (counter, slot) hash, exactly as ``_apply_force_accepted_tokens`` does. ``force_vec_req_kernel``: +one draw per request from its own key (its KV length at this verify), no shared counter. +""" + +import triton +import triton.language as tl + + +# One compiled variant for every batch and accept count: specializing them (== 1, divisible by 16) would compile +# mid-serving the first time a step reaches a new one, past the worker's warm-up. +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_kernel( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK: tl.constexpr, +): + c = tl.load(counter) + 1 + slots = tl.arange(0, BLOCK) + tl.store(counter, c) + idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask + live = slots < batch_size + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) + + +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_req_kernel( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK: tl.constexpr, +): + # One draw per request from its own key (its KV length at this verify): no shared counter, so a request's forced + # counts depend only on its own prompt and progress, not on what else the server ran before or beside it. + slots = tl.arange(0, BLOCK) + live = slots < batch_size + k = tl.load(keys + slots, mask=live, other=0).to(tl.int64) + idx = (k * key_stride) & pool_mask + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fullctx_decode.py new file mode 100644 index 000000000000..71666eb05872 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fullctx_decode.py @@ -0,0 +1,332 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 full-context FP8 decode attention on Rubin: ``fullctx_decode``. + +The dense layers (0-2) attend every token of their FP8 E4M3 (unscaled) paged KV cache. On SM 10.7 the stock call +(flashinfer trtllm-gen, ``minimax_m3_trtllm_gen_dense_decode``) caps its split at 16 cluster CTAs per (request, KV +head): at TP8 (8 query heads and 1 KV head per rank) one request runs on 16 CTAs of 212. + +This runs the same arithmetic -- FP8 Q (the same E4M3 rounding of q), FP8 K/V, P quantized to E4M3 per split -- on +the vendored PrimTS decode kernel (``attention/backends/prims_ts``, SM 10.7 through the sm_100f family) at a pinned +split-KV recipe, which fans one request out over many more CTAs. The split partials are merged by +``torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce`` (``cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/ +fullCtxReduce.cu``, head-count generic, used as is): unnormalized fp32 accumulation over the active splits, one final +division, one BF16 rounding, the merge form of stock trtllm-gen. Only the split changes against stock +(modeling's ``FULLCTX_DECODE_RECIPES``). + +A plan is compiled per (device, query rows per request, requests) by ``prebuild`` before the executor's warmup; the +forward only reads plans (``plan_for``). The plan is built from PrimTS's own config, compile and workspace helpers at +an explicit recipe; the PrimTS wrapper's planning policy is not used, so no PrimTS state is replaced. With +``own_reducer`` off the PrimTS reducer merges into FP16 and the result is cast to BF16. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Mapping, Optional, Tuple + +import torch + +# The TP8 per-rank geometry: 8 query heads over 1 KV head. +HEADS = 8 +KV_HEADS = 1 # capacity: fixed: the tp8 KV heads per rank +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim +PAGE = 128 +# The per-request K/V length PrimTS's config selection sizes its policy for (Q/KV tiles, launch mode); the pinned +# recipes were measured with it. The compiled capacity (``max_kv_len``) may be larger. +POLICY_KV_LEN = ( + 1 << 20 +) # capacity: fixed: the PrimTS policy's K/V length bound (the plans' compile shape) +HEADS_PER_KV = HEADS // KV_HEADS # capacity: fixed: the tp2 query heads per KV head + +_plans: Dict[Tuple[int, int, int, bool], "FullCtxPlan"] = {} +_lock = threading.Lock() + + +class FullCtxPlan: + """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" + + def __init__( + self, + device: torch.device, + qlen: int, + nreq: int, + max_kv_len: int, + recipe: Mapping, + own_reducer: bool, + pdl: bool, + ): + from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec + from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( + get_max_active_clusters_for_cluster_size, + make_decode_config, + ) + + self._tsdec = tsdec + self.qlen = qlen + self.nreq = nreq + self.pdl = pdl + self.recipe = dict(recipe) + splits = int(self.recipe["splits"]) + fp8 = tsdec._cutlass_dtype("float8_e4m3fn") + # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer + # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents + # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an + # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the + # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own + # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge + # (minimax_m3_sm107_tp2_fullctx_reduce) waits as its first instruction. + cfg = make_decode_config( + headdim=HEAD_DIM, + args={"use_pdl": pdl, **self.recipe.get("config", {})}, + seq_len_q=qlen, + seq_len_kv=min(max_kv_len, POLICY_KV_LEN), + batch_size=nreq, + num_heads_q=HEADS, + num_heads_kv=KV_HEADS, + q_dtype=fp8, + k_dtype=fp8, + v_dtype=fp8, + o_dtype=tsdec._cutlass_dtype("float16"), + qkv_layout="pagedKv", + num_tokens_per_page=PAGE, + storage_tokens_per_page=PAGE, + split_kv_mode="gmem_reduction_with_separate_kernel", + splits_kv=splits, + max_splits_kv=splits, + min_loop_iters_per_split=1, + mask_type="causal" if qlen > 1 else "dense", + auto_tuner=True, + split_kv=True, + ) + _check_layout(cfg, qlen, splits) + with torch.cuda.device(device): + max_active = get_max_active_clusters_for_cluster_size(1) + spec = tsdec._decode_launch_spec_from_config( + cfg, + batch_size=nreq, + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + seq_len_q=qlen, + max_active_clusters=int(max_active), + ) + compile_spec = tsdec._make_decode_compile_spec( + spec, + device_index=_index(device), + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + page_size=PAGE, + max_kv_len=max_kv_len, + seq_len_q=qlen, + q_dtype_key="float8_e4m3fn", + k_dtype_key="float8_e4m3fn", + v_dtype_key="float8_e4m3fn", + output_dtype_key="float16", + use_packed_q=False, + kv_prefix_mode="dynamic", + kv_lengths_mode="dynamic", + ) + self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) + layout = tsdec._make_decode_workspace_layout( + spec.scratch_shapes, + torch.float16, + use_separate_reduction_kernel=True, + use_split_kv=True, + ) + self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) + self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) + self.workspace.split_kv_counter.zero_() + self.workspace.cu_seqlens_q.zero_() + self.workspace.attention_sinks.zero_() + rows = HEADS_PER_KV * qlen + if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): + raise RuntimeError( + f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " + f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" + ) + self.own_reducer = own_reducer + self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) + self.tile_kv = int(cfg.tile_size_kv) + self.insts_kv = int(cfg.num_insts_kv) + self.causal = qlen > 1 + shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) + self.staging = torch.empty(shape, dtype=torch.float16, device=device) + + def run( + self, + q: torch.Tensor, + k_pages: torch.Tensor, + v_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + mxfp8: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> bool: + """Decodes into output; with mxfp8 (E4M3 rows, UE8M0 scales) and the TP8 merge, the merge also writes + the MXFP8 copy of output. Returns whether it did.""" + # The kernel takes a compact query; a column view of the fused QKV projection is copied once. + q = q.reshape(self.staging.shape).contiguous() + runtime = self._tsdec._prepare_decode_runtime_unchecked( + q, + (k_pages, v_pages), + None, + output_dtype=torch.float16, + bmm1_scale=sm_scale, + bmm2_scale=1.0, + out=self.staging, + ) + ws = self.workspace + self._main( + runtime.q, + runtime.k_cache, + runtime.v_cache, + runtime.k_sf_cache, + runtime.v_sf_cache, + runtime.out, + seq_lens, + ws.cu_seqlens_q, + block_table, + block_table, # Dense decode constexpr-elides the membership table. + ws.partial_o, + ws.partial_stats, + ws.split_kv_counter, + ws.attention_sinks, + runtime.num_physical_pages, + runtime.k_page_stride, + runtime.k_head_stride, + runtime.k_token_stride, + runtime.v_page_stride, + runtime.v_head_stride, + runtime.v_token_stride, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + if self.own_reducer: + reduce = _reduce_op() + args = ( + ws.partial_o, + ws.partial_stats, + seq_lens, + output, + self.qlen, + self.q_tokens_per_cta, + self.tile_kv, + self.insts_kv, + self.causal, + self.pdl, + ) + if mxfp8 is not None and reduce is _tp8_reduce_op(): + reduce(*args, mxfp8[0], mxfp8[1]) + return True + reduce(*args) + return False + self._primts_reducer( + runtime.out, + seq_lens, + ws.cu_seqlens_q, + ws.partial_o, + ws.partial_stats, + ws.attention_sinks, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + output.view(self.staging.shape).copy_(self.staging) + return False + + +def _check_layout(cfg, qlen: int, splits: int) -> None: + """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. + + Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is + ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding + window, attention sinks or packed Q. + """ + problems = [] + if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): + problems.append("not split-KV with the separate GMEM reducer") + if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: + problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") + if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: + problems.append("cluster-SMEM reduction or a persistent scheduler") + if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: + problems.append("packed Q, a sliding window or attention sinks") + if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): + problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") + if qlen > 1: + if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): + problems.append("multi-token Q without the grouped [token, head] tile layout") + elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): + problems.append( + f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" + ) + if problems: + raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + + +# The split-KV merge runs the TP8 copy of the BF16 reducer when the op library carries it (the same bits). +TP8_REDUCER = True + + +def _reduce_op(): + """The BF16 split-KV merge: the TP8 copy (one row per CTA, its dependent released right after its wait) when the + TP8 op library carries it, else the image's TP2 op; the same per-row arithmetic, so the same bits.""" + ops = torch.ops.trtllm + if TP8_REDUCER and hasattr(ops, "minimax_m3_sm107tp8_fullctx_reduce"): + return ops.minimax_m3_sm107tp8_fullctx_reduce + return ops.minimax_m3_sm107_tp2_fullctx_reduce + + +def _tp8_reduce_op(): + return getattr(torch.ops.trtllm, "minimax_m3_sm107tp8_fullctx_reduce", None) + + +def emits_mxfp8() -> bool: + """Whether the merge run uses can also write the o_proj MXFP8 hand-off (the TP8 copy of the merge).""" + return _tp8_reduce_op() is not None and _reduce_op() is _tp8_reduce_op() + + +def has_own_reducer() -> bool: + """Whether this build carries a BF16 split-KV merge op (the TP8 copy or the image's TP2 op).""" + ops = torch.ops.trtllm + return hasattr(ops, "minimax_m3_sm107tp8_fullctx_reduce") or hasattr( + ops, "minimax_m3_sm107_tp2_fullctx_reduce" + ) + + +def plan_for( + device: torch.device, qlen: int, nreq: int, own_reducer: bool = True +) -> Optional[FullCtxPlan]: + """The prebuilt plan for this shape, or None.""" + return _plans.get((_index(device), qlen, nreq, own_reducer)) + + +def prebuild( + device: torch.device, + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + pdl: bool, +) -> int: + """Compile the plan of every (query rows per request, requests) in ``recipes``. Call before warmup.""" + built = 0 + for (qlen, nreq), recipe in recipes.items(): + key = (_index(device), qlen, nreq, own_reducer) + with _lock: + if key not in _plans: + _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) + built += 1 + return built + + +def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: + """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" + return pool[0::subpages_per_slot], pool[1::subpages_per_slot] + + +def _index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_index_score.py new file mode 100644 index 000000000000..1467852d67b0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_index_score.py @@ -0,0 +1,692 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with four schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. +* ``prestage_old`` (with ``early_issue``, ``INDEX_SCORE_PRESTAGE_OLD``): when the CTA's first block holds only + tokens older than this step (block < (seq_lens hint - decode query length) / 128), its K tile is TMA-loaded into its + stage before the wait; Q (and a first block that is not all old) is still loaded after it. +* ``early_trigger`` (``INDEX_SCORE_SET``): every CTA releases the dependent block selector right after its + setup, so the selector launches and makes its pre-wait reads while this grid scores; the selector still waits for + this whole grid before it reads a score, so the scores and the selection are unchanged. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer``) through ``IndexScoreRunner`` +and called through this target's ``m3_index_score`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import ( + TRTLLM_ENABLE_PDL, + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) +from tensorrt_llm.logger import logger + +# With the early issue, a CTA whose first block holds only tokens older than this step also TMA-loads that block's +# index-K tile before the grid wait (bit-identical scores). +INDEX_SCORE_PRESTAGE_OLD = True +# Every CTA releases the dependent block selector right after its setup instead of at its exit (the selector's own +# grid wait still returns only once the whole scoring grid is complete, so the picks are unchanged), at every key. +INDEX_SCORE_SET = True + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + prestage_old: bool = False, + early_trigger: bool = False, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K + # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's + # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. + self.prestage_old = prestage_old and early_issue + # early_trigger: release the dependent block selector right after this CTA's setup (see the module docstring). + self.early_trigger = early_trigger + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # Without early_trigger the dependent block selector is released as the CTAs exit. + early_q = False + early_k = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + seqlen_hint = seq_lens[batch_id] + hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + if cutlass.const_expr(self.prestage_old): + # Stage the first K tile now when every token of the block predates this step. + if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: + early_k = True + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + if not early_k: + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # The length that bounds this grid's work is read after the wait (the + # pre-wait reads above are prefetch / early-issue hints only, clamped to + # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens + # patches every step, and num_blocks bounds an unmasked block_table load + # and an unmasked score store, so a stale length is an out-of-bounds write + # rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA; ``early_trigger`` is ``INDEX_SCORE_SET`` for every key of the process.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + prestage_old=early and INDEX_SCORE_PRESTAGE_OLD, + early_trigger=INDEX_SCORE_SET, + ) + if INDEX_SCORE_SET and not cls.kernel_cache: + logger.info( + "tp8 index scorer: the block selector is released after the scorer's setup (INDEX_SCORE_SET)" + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/phase_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/phase_producer.py new file mode 100644 index 000000000000..c8dfa3eee847 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/phase_producer.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode producer as a phase of the TP8 per-layer kernel (m3_layer): q / k / index-q / index-k RMSNorm + NeoX +partial RoPE, NVFP4 K / V and FP8 index-K cache insert, FP8 q / index-q out. + +Bit for bit ``torch.ops.trtllm.minimax_m3_producer_pw`` (cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu, +m3_producer_kernel<8, *>) at the TP8 geometry: 8 query heads, 1 KV head, 1 index head, head_dim 128, rotary_dim 64, +page 128. One work item = one decode token; the CTA's warp w (< 12) owns packed head w of it, lane l holds elements +[4l, 4l + 4) exactly as the C++ kernel, so every reduction (the RMS sum of squares, the RoPE pair exchange, the +NVFP4 block max) runs the same shuffle tree in the same order. Every floating-point operation is written as its PTX +instruction (mul.rn / add.rn / fma.rn / rsqrt.approx.ftz / rcp.approx.ftz / cvt) so the DSL compiler cannot contract +or reassociate differently from the library build (--use_fast_math). The places where that build contracted are +marked CONTRACT and follow the shipped SASS (msa/results/sassdump-*). + +Packed row layout per token: [8 q | k | v | index_q | index_k] x 128 BF16 (the QKV GEMV output rounded to BF16). +""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op + +HEADS_Q = 8 # capacity: fixed: the tp8 query heads per rank +PACKED_HEADS = HEADS_Q + 4 # q x 8, k, v, index_q, index_k +HEAD_DIM = 128 +ROTARY_DIM = 64 # capacity: fixed: the checkpoint's rotary dim +PAGE = 128 +ELEMS = HEAD_DIM // 32 # capacity: fixed: elements per lane (per lane) +K_HEAD, V_HEAD, IQ_HEAD, IK_HEAD = HEADS_Q, HEADS_Q + 1, HEADS_Q + 2, HEADS_Q + 3 + + +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- exact float arithmetic (no contraction) --------------------------------------------------------------------- +@dsl_user_op +def f_mul(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_add(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_fma(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_max(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "max.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_abs(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "abs.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_neg(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "neg.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_rsqrt(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rsqrt.approx.ftz.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_rcp(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def bf16_bits_to_f32(bits16, *, loc=None, ip=None): + """The float value of a BF16 given in the low 16 bits of an Int32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(bits16)], + "{\n\t.reg .b32 t;\n\tshl.b32 t, $1, 16;\n\tmov.b32 $0, t;\n\t}", + "=f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def round_bf16(a, *, loc=None, ip=None): + """__bfloat162float(__float2bfloat16_rn(a)).""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a)], + "{\n\t.reg .b16 h;\n\t.reg .b32 t;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b32 t, {0, h};\n\tmov.b32 $0, t;\n\t}", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e4m3x4_word(e0, e1, e2, e3, *, loc=None, ip=None): + """store_fp8x4: bytes [e0, e1, e2, e3] (cvt.rn.satfinite.e4m3x2.f32 puts its first source in the high byte).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(e0), cutlass.Float32(e1), cutlass.Float32(e2), cutlass.Float32(e3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", + "=r,f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e4m3_byte_and_value(x, *, loc=None, ip=None): + """(byte, float value) of __nv_fp8_e4m3(x): the encoded scale and static_cast of it.""" + byte = cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n\t.reg .b16 p;\n\t.reg .b32 w;\n\tcvt.rn.satfinite.e4m3x2.f32 p, 0f00000000, $1;\n\t" + "cvt.u32.u16 w, p;\n\tand.b32 $0, w, 0xFF;\n\t}", + "=r,f", + loc=loc, + ip=ip, + ) + ) + value = cutlass.Float32( + _asm( + _T.f32(), + [byte], + "{\n\t.reg .b16 p;\n\t.reg .b32 h2;\n\t.reg .b16 h;\n\tcvt.u16.u32 p, $1;\n\t" + "cvt.rn.f16x2.e4m3x2 h2, p;\n\tmov.b32 {h, _}, h2;\n\tcvt.f32.f16 $0, h;\n\t}", + "=f,r", + loc=loc, + ip=ip, + ) + ) + return byte, value + + +@dsl_user_op +def e2m1x4_halfword(q0, q1, q2, q3, *, loc=None, ip=None): + """The low 16 bits of fp32x8_to_e2m1 for q[0..3] (q[4..7] = 0): byte0 = (q1, q0), byte1 = (q3, q2).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(q0), cutlass.Float32(q1), cutlass.Float32(q2), cutlass.Float32(q3)], + "{\n\t.reg .b8 b0, b1;\n\t.reg .b32 w;\n\tcvt.rn.satfinite.e2m1x2.f32 b0, $2, $1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b1, $4, $3;\n\tmov.b32 w, {b0, b1, 0, 0};\n\tmov.b32 $0, w;\n\t}", + "=r,f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +# ---- global memory ----------------------------------------------------------------------------------------------- +@dsl_user_op +def ld_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm(_T.i64(), [cutlass.Int64(addr)], "ld.global.u64 $0, [$1];", "=l,l", True, loc, ip) + ) + + +@dsl_user_op +def ld_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm(_T.i32(), [cutlass.Int64(addr)], "ld.global.s32 $0, [$1];", "=r,l", True, loc, ip) + ) + + +@dsl_user_op +def ld_f32(addr, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Int64(addr)], "ld.global.f32 $0, [$1];", "=f,l", True, loc, ip) + ) + + +@dsl_user_op +def st_u32(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "st.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_u16(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "{\n\t.reg .b16 h;\n\tcvt.u16.u32 h, $1;\n\tst.global.u16 [$0], h;\n\t}", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_u8(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "{\n\t.reg .b16 h;\n\tcvt.u16.u32 h, $1;\n\tst.global.u8 [$0], h;\n\t}", + "l,r", + True, + loc, + ip, + ) + + +# ---- CONTRACT points (follow the shipped SASS; see the module docstring) ------------------------------------------- +def sumsq_pair(acc, x, y): + """sumSquares += x * x + y * y (C++: one pair of the packed row). CONTRACT: fma(y, y, x * x) then add.""" + return f_add(acc, f_fma(y, y, f_mul(x, x))) + + +def rms_reciprocal(sum_squares, eps): + """rsqrtf(sumSquares / 128 + eps). CONTRACT: the / 128 is an exact * 2^-7; fused with + eps into one fma.""" + return f_rsqrt(f_fma(sum_squares, cutlass.Float32(1.0 / HEAD_DIM), eps)) + + +class ProducerArgs: + """Raw device addresses (int64) + strides, as the C++ Params (producerPw.h) carries them.""" + + def __init__( + self, + packed, + q_out, + index_q_out, + index_k, + kv_data, + kv_scale, + out_cache_loc, + kv_quant_scale, + position_ids, + rotary_cos_sin, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + eps, + num_pages, + index_page_stride, + index_token_stride, + data_page_stride, + data_plane_stride, + data_token_stride, + scale_page_stride, + scale_plane_stride, + ): + self.__dict__.update(locals()) + del self.__dict__["self"] + + +@cute.jit +def phase_producer(tok, a: ProducerArgs): + """Producer work item ``tok``: warp w (< 12) of the calling CTA does packed head w of decode token tok. + The caller guarantees tok < rows and that warps >= 12 skip. Stores only (no CTA barrier inside).""" + tid = cute.arch.thread_idx()[0] + lane = tid % 32 + head = tid // 32 + in_off = (cutlass.Int64(tok) * PACKED_HEADS + cutlass.Int64(head)) * HEAD_DIM + cutlass.Int64( + lane * ELEMS + ) + pk = ld_u64(a.packed + in_off * 2) + x0 = bf16_bits_to_f32(cutlass.Int32(pk & cutlass.Int64(0xFFFF))) + x1 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(16)) & cutlass.Int64(0xFFFF))) + x2 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(32)) & cutlass.Int64(0xFFFF))) + x3 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(48)) & cutlass.Int64(0xFFFF))) + producer_head(tok, head, lane, x0, x1, x2, x3, a) + + +@cute.jit +def producer_head(tok, head, lane, x0, x1, x2, x3, a: ProducerArgs): + """Everything after the packed-row load for packed head ``head`` (warp-uniform; dynamic Int32 or Constexpr) of + decode token ``tok``, on the calling warp: lane ``lane`` holds the head's elements [4 lane, 4 lane + 4) as the + float values x0..x3 of the BF16 packed row (the QKV GEMV output rounded to BF16). Norm weights, KV quant scale, + position and RoPE loads, the RMS tree, RoPE, and the q / index-q / index-K / K-V NVFP4 stores. Shared by + phase_producer and the QKV GEMV owner epilogue (MK-A), so the producer math has one implementation. + + Data flow is branch-free (every load is in bounds for every head: V reads the q norm weight, a non-rotating lane + reads its token's RoPE row, and the results are then discarded with select); only the stores branch, on + warp-uniform conditions, so every shuffle runs on the full warp as in the C++ kernel.""" + l4 = cutlass.Int64(lane * ELEMS) + is_v = head == V_HEAD + zero = cutlass.Float32(0.0) + + # Norm weights (BF16): q / k / index-q / index-k weight by head (V: q's, unused). + w_base = cutlass.select_( + head == K_HEAD, + a.k_weight, + cutlass.select_( + head == IQ_HEAD, + a.index_q_weight, + cutlass.select_(head == IK_HEAD, a.index_k_weight, a.q_weight), + ), + ) + ww = ld_u64(w_base + l4 * 2) + w0 = bf16_bits_to_f32(cutlass.Int32(ww & cutlass.Int64(0xFFFF))) + w1 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(16)) & cutlass.Int64(0xFFFF))) + w2 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(32)) & cutlass.Int64(0xFFFF))) + w3 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(48)) & cutlass.Int64(0xFFFF))) + quant = ld_f32(a.kv_quant_scale + cutlass.select_(is_v, cutlass.Int64(8), cutlass.Int64(4))) + + # Position id, RoPE row. + pos = ld_s32(a.position_ids + cutlass.Int64(tok) * 4) + row = cutlass.Int64(pos) * ROTARY_DIM + coef = cutlass.Int64((lane * ELEMS) % (ROTARY_DIM // 2)) + c0 = ld_f32(a.rotary_cos_sin + (row + coef) * 4) + c1 = ld_f32(a.rotary_cos_sin + (row + coef + 1) * 4) + c2 = ld_f32(a.rotary_cos_sin + (row + coef + 2) * 4) + c3 = ld_f32(a.rotary_cos_sin + (row + coef + 3) * 4) + s0 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef) * 4) + s1 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 1) * 4) + s2 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 2) * 4) + s3 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 3) * 4) + + # RMSNorm (all heads compute it; V keeps x). + ss = sumsq_pair(zero, x0, x1) + ss = sumsq_pair(ss, x2, x3) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=16)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=8)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=4)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=2)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=1)) + r = rms_reciprocal(ss, a.eps) + one = cutlass.Float32(1.0) + n0 = f_mul(x0, f_mul(r, f_add(one, w0))) + n1 = f_mul(x1, f_mul(r, f_add(one, w1))) + n2 = f_mul(x2, f_mul(r, f_add(one, w2))) + n3 = f_mul(x3, f_mul(r, f_add(one, w3))) + # NeoX partial RoPE: lane l's dim 4l + i pairs with lane l ^ 8 (dim +-32) within the first 64 dims. + pair_off = (ROTARY_DIM // 2) // ELEMS + p0 = cute.arch.shuffle_sync_bfly(n0, offset=pair_off) + p1 = cute.arch.shuffle_sync_bfly(n1, offset=pair_off) + p2 = cute.arch.shuffle_sync_bfly(n2, offset=pair_off) + p3 = cute.arch.shuffle_sync_bfly(n3, offset=pair_off) + rot = lane * ELEMS < ROTARY_DIM + low = lane * ELEMS < ROTARY_DIM // 2 + p0 = cutlass.select_(low, f_neg(p0), p0) + p1 = cutlass.select_(low, f_neg(p1), p1) + p2 = cutlass.select_(low, f_neg(p2), p2) + p3 = cutlass.select_(low, f_neg(p3), p3) + r0 = cutlass.select_(rot, f_fma(p0, s0, f_mul(n0, c0)), n0) + r1 = cutlass.select_(rot, f_fma(p1, s1, f_mul(n1, c1)), n1) + r2 = cutlass.select_(rot, f_fma(p2, s2, f_mul(n2, c2)), n2) + r3 = cutlass.select_(rot, f_fma(p3, s3, f_mul(n3, c3)), n3) + e0 = cutlass.select_(is_v, x0, r0) + e1 = cutlass.select_(is_v, x1, r1) + e2 = cutlass.select_(is_v, x2, r2) + e3 = cutlass.select_(is_v, x3, r3) + + if head < HEADS_Q: + st_u32( + a.q_out + (cutlass.Int64(tok) * HEADS_Q + cutlass.Int64(head)) * HEAD_DIM + l4, + e4m3x4_word(e0, e1, e2, e3), + ) + if head == IQ_HEAD: + st_u32(a.index_q_out + cutlass.Int64(tok) * HEAD_DIM + l4, e4m3x4_word(e0, e1, e2, e3)) + if head >= K_HEAD: + if head != IQ_HEAD: + slot = ld_s32(a.out_cache_loc + cutlass.Int64(tok) * 4) + page = slot // PAGE + if slot >= 0: + if page < a.num_pages: + within = slot & (PAGE - 1) + if head == IK_HEAD: + st_u32( + a.index_k + + cutlass.Int64(page) * a.index_page_stride + + cutlass.Int64(within) * a.index_token_stride + + l4, + e4m3x4_word(e0, e1, e2, e3), + ) + else: + plane = cutlass.select_(is_v, cutlass.Int64(1), cutlass.Int64(0)) + col = lane // 4 + sw = cutlass.select_( + is_v, + cutlass.Int64((within // 4) * 32 + col * 4 + within % 4), + cutlass.Int64(within * 8 + col), + ) + s_off = ( + cutlass.Int64(page) * a.scale_page_stride + + plane * a.scale_plane_stride + + sw + ) + # quantize_fp4x4: BF16-rounded values, 16-element block max over the 4-lane group. + q0 = round_bf16(e0) + q1 = round_bf16(e1) + q2 = round_bf16(e2) + q3 = round_bf16(e3) + lmax = f_max( + f_max(f_max(f_max(zero, f_abs(q0)), f_abs(q1)), f_abs(q2)), f_abs(q3) + ) + lmax = f_max(lmax, cute.arch.shuffle_sync_bfly(lmax, offset=1)) + lmax = f_max(lmax, cute.arch.shuffle_sync_bfly(lmax, offset=2)) + sbyte, sval = e4m3_byte_and_value( + f_mul(quant, f_mul(lmax, f_rcp(cutlass.Float32(6.0)))) + ) + if lane % 4 == 0: + st_u8(a.kv_scale + s_off, sbyte) + outs = cutlass.select_(lmax != zero, f_rcp(f_mul(sval, f_rcp(quant))), zero) + hw = e2m1x4_halfword( + f_mul(q0, outs), f_mul(q1, outs), f_mul(q2, outs), f_mul(q3, outs) + ) + d_off = ( + cutlass.Int64(page) * a.data_page_stride + + plane * a.data_plane_stride + + cutlass.Int64(within) * a.data_token_stride + + cutlass.Int64(lane * 2) + ) + st_u16(a.kv_data + d_off, hw) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_locality.py new file mode 100644 index 000000000000..3e88c7fb1ce2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_locality.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Locality domains ("uGPU") of a multi-die GPU (Rubin VR200: 2 domains of 100 SMs + 12 remainder SMs) for the +streaming decode kernels of this target: the SM -> domain map, per-domain weight placement, and the CTA -> SM +assignment a persistent kernel uses to read only its own die's memory. + +A kernel launched with one CTA per SM reads ``%smid`` once and takes its logical block index from a host-built +permutation (:meth:`LocalityMap.logical_map`), so every per-CTA structure (work plan row, epochs, roles) follows its +die. Weights that are split by output rows into per-domain halves (:meth:`LocalityMap.split_rows`) live in one virtual +range whose halves are physically in the two domains. Placement only: the arithmetic is unchanged, so results are +bitwise those of the unplaced weights. + +The SM -> domain map is attested by the driver: the SMs are split by locality domain +(``cuDevSmResourceSplit`` with ``CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID``), and a ``%smid`` probe runs on a green +context of each split. Allocations use driver VMM with ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN`` and are checked +with ``CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL``; any failure raises (fail closed).""" + +from __future__ import annotations + +import threading +from typing import Dict, List, Sequence, Tuple + +import torch + +_LOCK = threading.Lock() +_MAPS: Dict[int, "LocalityMap"] = {} +_PROBE_SRC = r""" +extern "C" __global__ void smid_probe(int* out) { + if (threadIdx.x == 0) { + unsigned s; + asm volatile("mov.u32 %0, %%smid;" : "=r"(s)); + out[blockIdx.x] = (int)s; + } +} +""" + + +def _drv(): + from cuda.bindings import driver + + return driver + + +def _ck(res, what: str): + """cuda-python returns (CUresult, *values): raise on an error, return the values.""" + err = res[0] if isinstance(res, tuple) else res + if int(err) != 0: + raise RuntimeError(f"locality: {what} failed: {err}") + if isinstance(res, tuple): + return res[1] if len(res) == 2 else res[1:] + return None + + +class _Mapping: + """One virtual range backed by per-domain physical allocations, exposed to torch through + ``__cuda_array_interface__``. Lives as long as a tensor that views it (the tensor holds this object).""" + + def __init__(self, device: int, pieces: Sequence[Tuple[int, int]]): + """``pieces``: (bytes, domain) in address order; each a multiple of the allocation granularity.""" + d = _drv() + self.device = device + self.size = sum(n for n, _ in pieces) + self.va = _ck(d.cuMemAddressReserve(self.size, 0, 0, 0), "cuMemAddressReserve") + self.handles = [] + off = 0 + try: + for nbytes, dom in pieces: + prop = _alloc_prop(device, dom) + h = _ck(d.cuMemCreate(nbytes, prop, 0), f"cuMemCreate(domain {dom})") + self.handles.append((h, nbytes)) + _ck(d.cuMemMap(int(self.va) + off, nbytes, 0, h, 0), "cuMemMap") + off += nbytes + acc = d.CUmemAccessDesc() + acc.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + acc.location.id = device + acc.flags = d.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _ck(d.cuMemSetAccess(self.va, self.size, [acc], 1), "cuMemSetAccess") + except Exception: + self.close() + raise + self.__cuda_array_interface__ = { + "shape": (self.size,), + "typestr": "|u1", + "data": (int(self.va), False), + "version": 3, + "strides": None, + "stream": None, + } + + def close(self): + d = _drv() + off = 0 + for h, n in self.handles: + d.cuMemUnmap(int(self.va) + off, n) + d.cuMemRelease(h) + off += n + self.handles = [] + if self.va is not None: + d.cuMemAddressFree(self.va, self.size) + self.va = None + + def __del__(self): + try: + self.close() + except Exception: + pass + + +_LIBCUDA = None + + +def _libcuda(): + global _LIBCUDA + if _LIBCUDA is None: + import ctypes + + _LIBCUDA = ctypes.CDLL("libcuda.so.1") + return _LIBCUDA + + +def _alloc_prop(device: int, domain: int): + d = _drv() + prop = d.CUmemAllocationProp() + prop.type = d.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.requestedHandleTypes = d.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE + prop.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN + prop.location.localized.deviceId = device + prop.location.localized.localityDomainId = domain + return prop + + +class LocalityMap: + """The locality domains of one device. Build with :meth:`get` before CUDA-graph capture.""" + + def __init__(self, device: int): + d = _drv() + _ck(d.cuInit(0), "cuInit") + dev = _ck(d.cuDeviceGet(device), "cuDeviceGet") + attr = d.CUdevice_attribute + self.device = device + self.n_domains = int( + _ck( + d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, dev), + "domain count", + ) + ) + self.sms_per_domain = int( + _ck( + d.cuDeviceGetAttribute( + attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_MULTIPROCESSOR_COUNT, dev + ), + "SMs per domain", + ) + ) + self.num_sms = int( + _ck(d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev), "SMs") + ) + if self.n_domains != 2: + raise RuntimeError( + f"locality: {self.n_domains} locality domains (this layer handles 2)" + ) + self._domain_sms = self._probe(d, dev) + dom = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for s in sms: + if dom[s] != -1: + raise RuntimeError(f"locality: SM {s} probed in two domains") + dom[s] = k + self.sm_domain: List[int] = dom + self.remainder: List[int] = [s for s in range(self.num_sms) if dom[s] < 0] + rt = _alloc_prop(device, 0) + self.granularity = int( + _ck( + d.cuMemGetAllocationGranularity( + rt, d.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM + ), + "granularity", + ) + ) + + @classmethod + def get(cls, device) -> "LocalityMap": + idx = torch.device(device).index if not isinstance(device, int) else device + idx = torch.cuda.current_device() if idx is None else idx + with _LOCK: + m = _MAPS.get(idx) + if m is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "LocalityMap probes the device: build it before CUDA-graph capture" + ) + m = _MAPS[idx] = cls(idx) + return m + + def sms(self, domain: int) -> List[int]: + return list(self._domain_sms[domain]) + + @property + def sm_table(self) -> torch.Tensor: + """int32 [num_sms] on the device: (domain << 16) | rank within the domain's sorted SMs; -1 = remainder SM. + For kernels that pick work by %smid themselves.""" + t = getattr(self, "_sm_table", None) + if t is None: + v = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for i, s in enumerate(sms): + v[s] = (k << 16) | i + t = self._sm_table = torch.tensor(v, dtype=torch.int32, device=f"cuda:{self.device}") + return t + + def _probe(self, d, dev) -> List[List[int]]: + """Split the SMs by locality domain (strict, no backfill) and record %smid of a probe kernel launched on a + green context of each part: the driver's own SM -> domain assignment.""" + from cuda.bindings import nvrtc + + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + # the probe module, compiled for this device + major = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev + ), + "cc", + ) + minor = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev + ), + "cc", + ) + prog = _ck( + nvrtc.nvrtcCreateProgram(_PROBE_SRC.encode(), b"smid_probe.cu", 0, [], []), + "nvrtcCreateProgram", + ) + opts = [f"--gpu-architecture=sm_{major}{minor}".encode()] + err = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)[0] + if int(err) != 0: + raise RuntimeError(f"locality: probe compile failed: {err}") + n = _ck(nvrtc.nvrtcGetCUBINSize(prog), "cubin size") + cubin = b" " * n + _ck(nvrtc.nvrtcGetCUBIN(prog, cubin), "cubin") + out: List[List[int]] = [] + nb = 4 * self.sms_per_domain # more blocks than SMs: every SM of the part shows up + torch.cuda.set_device(self.device) + buf = torch.full((nb,), -1, dtype=torch.int32, device=f"cuda:{self.device}") + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate(desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM), + "green ctx", + ) + st = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + ctx = _ck(d.cuCtxFromGreenCtx(g), "ctx from green ctx") + _ck(d.cuCtxPushCurrent(ctx), "push ctx") + try: + mod = _ck(d.cuModuleLoadData(cubin), "probe module") + fn = _ck(d.cuModuleGetFunction(mod, b"smid_probe"), "probe fn") + buf.fill_(-1) + torch.cuda.synchronize(self.device) + import ctypes + + arg = ctypes.c_void_p(buf.data_ptr()) + _ck( + d.cuLaunchKernel( + fn, nb, 1, 1, 32, 1, 1, 0, st, ((arg,), (ctypes.c_void_p,)), 0 + ), + "probe launch", + ) + _ck(d.cuStreamSynchronize(st), "probe sync") + sms = sorted(set(int(x) for x in buf.tolist() if x >= 0)) + d.cuModuleUnload(mod) + finally: + d.cuCtxPopCurrent() + d.cuStreamDestroy(st) + d.cuGreenCtxDestroy(g) + if len(sms) != self.sms_per_domain: + raise RuntimeError( + f"locality: domain {k} probed {len(sms)} SMs, expected {self.sms_per_domain}" + ) + out.append(sms) + return out + + def logical_map(self, per_domain: int, extra: int) -> Tuple[torch.Tensor, List[int]]: + """The SM -> logical block index permutation for a kernel of ``2 * per_domain + extra`` CTAs, one per SM: + logical blocks [d * per_domain, (d + 1) * per_domain) run on domain d's first ``per_domain`` SMs; logical + blocks 2 * per_domain .. run on the rest (the leftover domain SMs, alternating domains, then the remainder + SMs). Returns (int32 [num_sms] on the device, -1 = no block; the domain of each extra block, -1 remainder).""" + if 2 * per_domain + extra != self.num_sms: + raise ValueError( + f"locality: {2 * per_domain} + {extra} CTAs must cover the {self.num_sms} SMs" + ) + if per_domain > self.sms_per_domain: + raise ValueError(f"locality: {per_domain} CTAs per domain > {self.sms_per_domain} SMs") + lmap = [-1] * self.num_sms + for k in range(2): + for i, s in enumerate(self._domain_sms[k][:per_domain]): + lmap[s] = k * per_domain + i + left = [self._domain_sms[k][per_domain:] for k in range(2)] + rest = [x for pair in zip(*left) for x in pair] + self.remainder + ext_dom = [] + for i, s in enumerate(rest): + lmap[s] = 2 * per_domain + i + ext_dom.append(self.sm_domain[s]) + assert sorted(lmap) == list(range(self.num_sms)) + return torch.tensor(lmap, dtype=torch.int32, device=f"cuda:{self.device}"), ext_dom + + def domain_streams(self) -> List[torch.cuda.ExternalStream]: + """One non-blocking stream per domain on a green context holding exactly that domain's SMs (kept alive with + this map): work launched on stream d runs only on domain d's SMs. Fork / join with events from the caller's + stream; graph capture follows the forks. Built once, before capture.""" + st = getattr(self, "_streams", None) + if st is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.domain_streams: build before CUDA-graph capture") + d = _drv() + dev = _ck(d.cuDeviceGet(self.device), "cuDeviceGet") + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + self._green = [] + st = [] + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate( + desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM + ), + "green", + ) + s_ = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + self._green.append((g, s_)) + st.append(torch.cuda.ExternalStream(int(s_), device=f"cuda:{self.device}")) + self._streams = st + return st + + def domain_of(self, ptr: int) -> int: + """The locality domain ordinal of a device address (-1: not placed in one domain). Through libcuda + directly: cuda-python's attribute helper does not know this attribute.""" + import ctypes + + d = _drv() + lib = _libcuda() + out = ctypes.c_int(-9) + attr = int(d.CUpointer_attribute.CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL) + err = lib.cuPointerGetAttribute(ctypes.byref(out), ctypes.c_int(attr), ctypes.c_uint64(ptr)) + if err != 0: + raise RuntimeError( + f"locality: cuPointerGetAttribute(LOCALITY_DOMAIN_ORDINAL) failed: {err}" + ) + return int(out.value) + + def place_split(self, t: torch.Tensor, split_row: int) -> torch.Tensor: + """Split-mapped placement: a copy of ``t`` (contiguous, rows along dim 0) with the SAME shape and layout, in one + virtual range whose rows [0, split_row) are physically in domain 0 and rows [split_row, N) in domain 1. The + split must fall on an allocation granule (rows [0, split_row) a multiple of 2 MiB on VR200); the tail is padded + to a granule in domain 1. Stock consumers read it as a plain tensor; no extra memory once ``t`` is dropped.""" + if not t.is_contiguous(): + raise ValueError("locality.place_split: needs a contiguous tensor") + row_bytes = t[0].numel() * t.element_size() if t.dim() > 1 else t.element_size() + head = split_row * row_bytes + total = t.numel() * t.element_size() + if not 0 < head < total or head % self.granularity: + raise ValueError( + f"locality.place_split: rows [0, {split_row}) = {head} B must be a granule multiple " + f"({self.granularity} B) inside the {total} B tensor" + ) + tail = -(-(total - head) // self.granularity) * self.granularity + m = _Mapping(self.device, [(head, 0), (tail, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + out = flat[:total].view(t.dtype).view(t.shape) + out.copy_(t) + out._locality_mapping = m + for q, k in ( + (out.data_ptr(), 0), + (out.data_ptr() + head - 1, 0), + (out.data_ptr() + head, 1), + (out.data_ptr() + total - 1, 1), + ): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.place_split: byte {q:#x} is not in domain {k}") + return out + + def replace_param_(self, param: torch.nn.Parameter, split_row: int) -> torch.nn.Parameter: + """Re-place a parameter in place with :meth:`place_split` (swaps ``param.data``; the old block returns to the + caching allocator, so call ``torch.cuda.empty_cache()`` between large re-placements).""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.replace_param_: not during CUDA-graph capture") + param.data = self.place_split(param.data, split_row) + return param + + def split_rows(self, t: torch.Tensor, dim: int = 1) -> torch.Tensor: + """Place ``t`` (contiguous, 1-byte elements, rows along ``dim`` split into two equal halves) as + [2, *t.shape[:dim], t.shape[dim] // 2, *t.shape[dim + 1:]]: half k (rows [k h, (k + 1) h) of every leading + index) in domain k's memory, one virtual range. Each half's byte size must be a multiple of the allocation + granularity (2 MiB on VR200). The copy is exact; the caller drops ``t`` afterwards.""" + if t.element_size() != 1 or not t.is_contiguous(): + raise ValueError("locality.split_rows: needs a contiguous 1-byte tensor") + n = t.shape[dim] + if n % 2: + raise ValueError(f"locality.split_rows: {n} rows don't split in halves") + half_bytes = t.numel() // 2 + if half_bytes % self.granularity: + raise ValueError( + f"locality.split_rows: a half is {half_bytes} B, not a multiple of the {self.granularity} B granularity" + ) + m = _Mapping(self.device, [(half_bytes, 0), (half_bytes, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + shape = (2,) + tuple(t.shape[:dim]) + (n // 2,) + tuple(t.shape[dim + 1 :]) + out = flat.view(torch.uint8).view(shape) + src = t.view(torch.uint8) + for k in range(2): + out[k].copy_(src.narrow(dim, k * (n // 2), n // 2)) + out._locality_mapping = m # the tensor keeps its mapping alive + for k in range(2): + p = out[k].data_ptr() + for q in (p, p + half_bytes - 1): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.split_rows: byte {q:#x} is not in domain {k}") + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_lowm_splitk.py new file mode 100644 index 000000000000..685797a96700 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_lowm_splitk.py @@ -0,0 +1,1592 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +This target's copy of FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the Eagle3 +drafter's low-M linears, modeling's ``DRAFT_LOWM_OWN_SPLITK``): the same kernel and the same split order at every +tactic, bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's +final cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, +the SwiGLU prologue that follow. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 # capacity: fixed: a kernel tile constant + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 # capacity: fixed: an alignment + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 # capacity: fixed: pipeline stages +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 # capacity: fixed: pipeline stages + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= _MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES + base = SplitKTactic(mma_m, mma_n, split_k, min_stages) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(min_stages, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", + "run_splitk_dense_swiglu", + "compile_swiglu_push", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _silu_mul_f32(g, u, *, loc=None, ip=None): + """silu(g) * (u + 0) in FP32 for one gate / up pair: the instruction sequence of the Triton silu_and_mul kernel at + its defaults (no limit, alpha 1, beta 0), so the bf16 rounding after it is bit for bit the Triton activation.""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [ + cutlass.Float32(g).ir_value(loc=loc, ip=ip), + cutlass.Float32(u).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .f32 t, e, s, p, q, one, zero;\n\t" + "mul.rn.f32 t, $1, 0fBFB8AA3B;\n\tex2.approx.f32 e, t;\n\t" + "mov.f32 one, 0f3F800000;\n\tadd.rn.f32 e, e, one;\n\tdiv.full.f32 s, one, e;\n\t" + "mul.rn.f32 p, $1, s;\n\tmov.f32 zero, 0f00000000;\n\tadd.rn.f32 q, $2, zero;\n\t" + "mul.rn.f32 $0, q, p;\n\t}", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _ld_bf16_f32(addr, *, loc=None, ip=None): + """The bf16 at global byte address ``addr``, widened to FP32 (exact).""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h;\n\tld.global.b16 h, [$1];\n\tcvt.f32.bf16 $0, h;\n\t}", + "=f,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_f32_bf16(addr, value, *, loc=None, ip=None) -> None: + """``value`` rounded to bf16 (round to nearest even) at global byte address ``addr``.""" + llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Float32(value).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b16 h;\n\tcvt.rn.bf16.f32 h, $1;\n\tst.global.b16 [$0], h;\n\t}", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None) -> None: + """Order this thread's generic global stores before later async-proxy (TMA) reads of them.""" + llvm.inline_asm( + None, + [], + "fence.proxy.async.global;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 # capacity: fixed: a split index, not a size + + +# The drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH): the exchange words of +# kernels/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.s32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + llvm.inline_asm( + None, + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + b_swiglu: bool = False, + ar_push: bool = False, + ) -> None: + self.ar_push = ar_push + self.acc_dtype = cutlass.Float32 + # The B operand (public A, the activation) is silu(gate) * up of a [M, 2K] gate | up tensor: the B DMA warp + # writes its CTA's K-slice of it to the B tensor (a scratch) before its unchanged TMA loop reads it back. + self.b_swiglu = b_swiglu + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + g: cute.Tensor, + u: cute.Tensor, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64 = cutlass.Int64(0), + ar_flags: cutlass.Int64 = cutlass.Int64(0), + ar_rank: Int32 = Int32(0), + ar_world: Int32 = Int32(1), + ar_maxt: Int32 = Int32(8), + ar_tokens: Int32 = Int32(0), + ar_n: Int32 = Int32(0), + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel( + a, + b, + c, + bias, + g, + u, + gu_ld, + b_ld, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + mG: cute.Tensor, # b_swiglu: the gate half (Gemm_N, Gemm_K, Gemm_L); dead otherwise + mU: cute.Tensor, # b_swiglu: the up half, same layout; dead otherwise + gu_ld: cutlass.Int32, # b_swiglu: elements between gate (and up) rows + b_ld: cutlass.Int32, # b_swiglu: elements between B rows + ar_ws: cutlass.Int64, # ar_push: the boundary all-reduce's multicast exchange words (int32) + ar_flags: cutlass.Int64, # ar_push: its per-(token, cluster CTA) call flags (parity = flag & 1) + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, # the bucket's tokens + ar_tokens: Int32, # rows pushed (public M) + ar_n: Int32, # public N (the hidden size): the exchange row width + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + # ar_push: the owner's BF16 tile (kernel-M = output column, kernel-N = token), M-major. + sD = cute_ext.allocate( + mC.element_type, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=16, + ) + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + cutlass.Int64(0), + cutlass.Int64(0), + cutlass.Int64(0), + Int32(0), + Int32(0), + Int32(0), + Int32(0), + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + mG.iterator.toint(), + mU.iterator.toint(), + mB.iterator.toint(), + gu_ld, + b_ld, + n_idx * Int32(self.cta_n), + Int32(cute.size(mB, mode=[0])) - n_idx * Int32(self.cta_n), + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sD, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + @cute.experimental.jit + def _swiglu_to_b( + self, + gate_base: cutlass.Int64, + up_base: cutlass.Int64, + out_base: cutlass.Int64, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + row0: cutlass.Int32, + rows_valid: cutlass.Int32, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + """This CTA's K-slice of the activation, silu(gate) * up rounded to bf16, into the B tensor's valid rows (TMA + fills the rest with zeros), through raw global addresses (the B tensor is in the TMA chain); then the stores + are made visible to the TMA reads that follow.""" + lane = cute.arch.lane_idx() + per_tile = cutlass.const_expr(self.cta_n * self.cta_k) + for k_tile in cutlass.range(k_tile_count, unroll=1): + k0 = (k_tile_start + k_tile) * Int32(self.cta_k) + for j in cutlass.range_constexpr(per_tile // 32): + idx = Int32(j * 32) + lane + row = idx // Int32(self.cta_k) + col = k0 + idx % Int32(self.cta_k) + if row < rows_valid: + src = ( + cutlass.Int64(row0 + row) * cutlass.Int64(gu_ld) + cutlass.Int64(col) + ) * cutlass.Int64(2) + dst = ( + cutlass.Int64(row0 + row) * cutlass.Int64(b_ld) + cutlass.Int64(col) + ) * cutlass.Int64(2) + act = _silu_mul_f32(_ld_bf16_f32(gate_base + src), _ld_bf16_f32(up_base + src)) + _st_f32_bf16(out_base + dst, act) + _fence_proxy_async_global() + cute.arch.sync_warp() + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + gate_base: cutlass.Int64, + up_base: cutlass.Int64, + out_base: cutlass.Int64, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + row0: cutlass.Int32, + rows_valid: cutlass.Int32, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not is_a and self.b_swiglu): + self._swiglu_to_b( + gate_base, + up_base, + out_base, + gu_ld, + b_ld, + row0, + rows_valid, + k_tile_start, + k_tile_count, + ) + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl and not self.ar_push): + # ar_push: the dependent is a prepushed all-reduce, which skips its grid wait: it is released only after + # this grid's own grid wait and its pushes, in the owner's _store_tile (the prepushed >= 2 contract). + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sD: cute.Tensor, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + self._store_tile(thr_t2r, rAcc, rD, gD_epi, c_dtype) + if cutlass.const_expr(self.ar_push): + self._push_tile( + thr_t2r, + rD, + sD, + epi_tid, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + @cute.experimental.jit + def _store_tile(self, thr_t2r, rAcc, rD, gD_epi, c_dtype: cutlass.Constexpr): + """The owner rank's reduced FP32 tile: one cast to ``c_dtype`` and the store to global memory.""" + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + @cute.experimental.jit + def _push_tile( + self, + thr_t2r, + rD, + sD, + epi_tid: cutlass.Int32, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output + would store them (kernels/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" + cute_ext.partition_and_copy(thr_t2r, rD, sD) + cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + pairs = self.cta_m // 2 + col0 = bidx * self.cta_m + half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) + for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): + word = epi_tid + w0 + t = word // pairs + p = word % pairs + if t < ar_tokens: + n = col0 + 2 * p + par = _ld_volatile_s32( + ar_flags + cutlass.Int64(4) * cutlass.Int64(t * 8 + n // (ar_n // 8)) + ) & Int32(1) + off = ( + cutlass.Int64(par) * half + + cutlass.Int64(t * ar_world + ar_rank) * cutlass.Int64(ar_n // 2) + + cutlass.Int64(n // 2) + ) + _st_global_b32( + ar_ws + off * cutlass.Int64(4), + _bf16x2_word( + sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + ), + ) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +import torch as _torch # noqa: E402 (as in FlashInfer) + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + b_sel = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + b_sel, + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + b_sel, + b_sel, + Int32(0), + Int32(0), + stream, + ) + + +@cute.experimental.jit +def _bmm_swiglu( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + g: cute.Tensor, + u: cute.Tensor, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + stream: _cuda.CUstream, +): + """_bmm_no_bias with the B operand built in the kernel from g / u (laid out as b, rows gu_ld apart; B rows b_ld + apart).""" + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + cute.make_tensor(g.iterator, cute.select(g.layout, mode=[2, 1, 0])), + cute.make_tensor(u.iterator, cute.select(u.layout, mode=[2, 1, 0])), + gu_ld, + b_ld, + stream, + ) + + +@cute.experimental.jit +def _bmm_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ar_tokens: Int32, + ar_n: Int32, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + Int32(0), + Int32(0), + stream, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + +@cute.experimental.jit +def _bmm_swiglu_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + g: cute.Tensor, + u: cute.Tensor, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ar_tokens: Int32, + ar_n: Int32, +): + """_bmm_swiglu that also pushes the output rows (_bmm_push's exchange).""" + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + cute.make_tensor(g.iterator, cute.select(g.layout, mode=[2, 1, 0])), + cute.make_tensor(u.iterator, cute.select(u.layout, mode=[2, 1, 0])), + gu_ld, + b_ld, + stream, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + b_sel = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + b_sel, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + b_sel, + b_sel, + Int32(0), + Int32(0), + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], + b_swiglu: bool = False, +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + b_swiglu=b_swiglu, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if b_swiglu: + # The gate / up halves are laid out as the B operand (public A) is. + gu = _make_compile_repr_tensors(dtype, False, *leading_dims)[0] + compiled = cute_ext.compile( + _bmm_swiglu, kernel, *compile_tensors[:3], gu, gu, Int32(2), Int32(1), stream + ) + elif has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any(tensor.dtype != a.dtype for tensor in tensors): + raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a, + b, + bias, + out, + pdl: bool, + tactic: SplitKTactic, +): + """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +@functools.cache +def _get_compiled_swiglu_push( + dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] +): + """The SwiGLU B load and the ar push in one kernel (the drafter down projection at a pushed boundary).""" + kernel = SplitKDenseGemmKernel( + tactic=tactic, use_pdl=use_pdl, has_bias=False, b_swiglu=True, ar_push=True + ) + compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) + gu = _make_compile_repr_tensors(dtype, False, *leading_dims)[0] + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + z64, z32 = cutlass.Int64(0), Int32(0) + return cute_ext.compile( + _bmm_swiglu_push, kernel, *compile_tensors[:3], gu, gu, Int32(2), Int32(1), stream, + z64, z64, z32, Int32(1), Int32(8), z32, Int32(6144), + ) # fmt: skip + + +def compile_swiglu_push(scratch, b, out, pdl: bool, tactic) -> None: + """Compile (only) the SwiGLU + push variant of ``run_splitk_dense_swiglu`` for these operands' layouts.""" + cute_tensors = _to_cute_swap(scratch, b, out, None) + _get_compiled_swiglu_push(scratch.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) + + +def run_splitk_dense_swiglu( + gate_up, b, out, pdl: bool, tactic: SplitKTactic, scratch, ar: tuple = None +): + """``out = silu_and_mul(gate_up) @ b`` in one launch: ``gate_up`` [M, 2K] (gate then up, contiguous rows) is + activated by the B DMA warp (the Triton silu_and_mul arithmetic at its defaults, bit for bit) into ``scratch`` + [>= M, K], which the unchanged GEMM then reads; the output is bit for bit ``run_splitk_dense(silu_and_mul(gate_up), + b, None, out, pdl, tactic)``. ``scratch`` must not be read or written by anything else while the kernel runs. With + ``ar`` (``run_splitk_dense_push``'s tuple) the rows are also pushed into the drafter boundary all-reduce's + exchange, from the same kernel.""" + if gate_up.ndim != 2 or gate_up.shape[1] % 2 or not gate_up.is_contiguous(): + raise ValueError("gate_up must be a contiguous [M, 2K] matrix") + m, k = gate_up.shape[0], gate_up.shape[1] // 2 + a = scratch[:m] + if a.shape != (m, k) or not a.is_contiguous(): + raise ValueError(f"scratch must be a contiguous [>= {m}, {k}] matrix") + gate, up = gate_up[:, :k], gate_up[:, k:] + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + if any(t.data_ptr() % 32 for t in (gate, up)) or gate.dtype != a.dtype: + raise ValueError("gate / up halves must be 32-byte aligned and of the GEMM dtype") + cute_tensors = _to_cute_swap(a, b, out, None) + g_swap = gate.unsqueeze(0).transpose(-2, -1) + u_swap = up.unsqueeze(0).transpose(-2, -1) + lead = cute_tensors[4][1] + if _detect_leading_dim(g_swap) != lead or _detect_leading_dim(u_swap) != lead: + raise ValueError("gate / up must share the activation's K-major layout") + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + operands = ( + *cute_tensors[:3], + _from_dlpack_dynamic(g_swap, lead), + _from_dlpack_dynamic(u_swap, lead), + Int32(gate_up.stride(0)), + Int32(a.stride(0)), + stream, + ) + if ar is None: + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=False, + leading_dims=cute_tensors[4], + b_swiglu=True, + ) + compiled(*operands) + return out + compiled = _get_compiled_swiglu_push(a.dtype, tactic, pdl, cute_tensors[4]) + ws, flags, rank, world, maxt = ar + compiled( + *operands, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + Int32(out.shape[0]), + Int32(out.shape[1]), + ) + return out + + +@functools.cache +def _get_compiled_push( + dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] +): + kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, ar_push=True) + compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + z64, z32 = cutlass.Int64(0), Int32(0) + return cute_ext.compile( + _bmm_push, + kernel, + *compile_tensors[:3], + stream, + z64, + z64, + z32, + Int32(1), + Int32(8), + z32, + Int32(6144), + ) + + +def compile_push(a, b, out, pdl: bool, tactic) -> None: + """Compile (only) the push variant of ``run_splitk_dense_push`` for these operands' shapes and layouts.""" + cute_tensors = _to_cute_swap(a, b, out, None) + _get_compiled_push(a.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) + + +def run_splitk_dense_push(a, b, out, pdl: bool, tactic: SplitKTactic, ar: tuple): + """``run_splitk_dense`` without bias that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange; + ``ar`` = (multicast words address, flags address, rank, world, bucket tokens) of the 8-token bucket.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_push(a.dtype, tactic, pdl, cute_tensors[4]) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + ws, flags, rank, world, maxt = ar + compiled( + *cute_tensors[:3], + stream, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + Int32(out.shape[0]), + Int32(out.shape[1]), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_mc_allreduce.py new file mode 100644 index 000000000000..bb87d8ead5bd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_mc_allreduce.py @@ -0,0 +1,439 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-latency TP all-reduce + residual + RMSNorm for decode (M <= 8 tokens) over a multicast mapping. + + partial = bf16(x + x2) (optional second input, e.g. shared + routed experts) + updated = bf16(residual + sum_r partial_r) (ranks summed in FP32, rank order) + normed = bf16(updated * rsqrt(mean(updated^2) + eps) * w) + +One cluster of CLUSTER CTAs per token; CTA c owns columns [c * H / CLUSTER, ...), one 16-byte chunk (8 BF16) +per thread. Each CTA pushes its slice of this rank's partial through the multicast mapping (one store lands in +every rank's buffer), releases dependents, polls its own buffer until all ranks' slices arrived, sums them, adds +the residual and reduces the RMS statistic across the cluster through DSMEM. + +Exchange protocol as in the K3 sandwich: int32 words, 0x80000000 = empty (pushed -0.0 halves become +0.0), each +CTA keeps a call counter whose parity selects the buffer half, readers re-arm the words they read. All ranks of +the TP group must issue the same calls in the same stream order. + +The workspace holds two exchange sets. With prepushed 3 two producers pushed this rank's x (set A) and x2 (set B), +e.g. the routed experts' finalize and the shared expert's down projection, and the kernel forms bf16(x + x2) per +rank from the two slots, exactly as the unpushed path does. + +Producers that push (the o_proj / shared-expert GEMV ``ar`` tuple, the routed-expert engine's ``arp`` descriptor) are +pointed at a bucket's exchange by its ``mc`` address, ``flags_for(8)`` counters, ``rank``, ``world_size`` and +``max_tokens``; ``get_workspace`` and ``_bucket`` select it. + +The kernel is ``torch.ops.trtllm.minimax_m3_mc_allreduce_norm``, reached through the ``ops/comm`` entries +``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``; this module owns the workspaces, the bucket and +epilogue-option choice and the MXFP8 publication. +""" + +import functools +from dataclasses import dataclass +from typing import Any + +import torch + +from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op +from ..ops.comm.m3_mc_finalize_allreduce_norm import ( + m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, +) +from . import m3_report + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +BUCKETS = (8, 16, 48) # token capacity of each workspace (grid = bucket * cluster CTAs) +EMPTY_WORD = -2147483648 + +# Attention / MoE boundary epilogue options (each bit-identical): 1 residual + norm weight loaded at entry (before +# the poll); 2 cluster RMS exchange by st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated +# right after the poll (no trailing __syncthreads); 8 max-smem carveout so the next kernel's CTAs co-reside early. +# The 8-token bucket, by token count 1..8 (GB300 TP4 tuning). +AR_OPTS_BY_M = (7, 7, 15, 15, 15, 15, 7, 7) +# The 16- and 48-token buckets. +AR_OPTS_BY_BUCKET = {16: 15, 48: 15} +# The MoE-finalize all-reduce, per bucket. +AR_FIN_OPTS_BY_BUCKET = {8: 7, 16: 15, 48: 15} + + +# Option 2 is also the safe ordering of the cluster RMS exchange: its stores into peer CTAs' shared memory follow a +# cluster barrier arrive / wait pair. Without it they run before any cluster barrier, when a peer need not have +# started, so every call keeps it. +_OPT_CLUSTER_BARRIER = 2 + + +def _ar_opts(m: int, bucket: int, finalize: bool = False) -> int: + """Epilogue options for an m-row call on the given bucket; every option is bit-identical. A bucket missing from + the tables gets 0, and every result keeps _OPT_CLUSTER_BARRIER.""" + if finalize: + opts = AR_FIN_OPTS_BY_BUCKET.get(bucket, 0) + elif bucket == 8: + opts = AR_OPTS_BY_M[m - 1] + else: + opts = AR_OPTS_BY_BUCKET.get(bucket, 0) + return opts | _OPT_CLUSTER_BARRIER + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags + + +class McBuckets: + """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" + + def __init__(self, workspaces): + self.workspaces = workspaces + self.hidden = workspaces[0].hidden + + def for_tokens(self, m: int): + for ws in self.workspaces: + if m <= ws.max_tokens: + return ws + return None + + +def _bucket(workspace, m): + return workspace.for_tokens(m) if isinstance(workspace, McBuckets) else workspace + + +def decline_reason(x, residual, workspace, cluster=8, x2=None) -> str: + """Why ``m3_mc_allreduce_norm`` returned None for these operands (for the dispatch's miss message).""" + if workspace is None: + return "no multicast workspace: it is created collectively on the first eager decode step" + bucket = _bucket(workspace, x.shape[0]) + if bucket is None: + return f"no workspace bucket holds {x.shape[0]} tokens (buckets {BUCKETS})" + if bucket.flags.get(cluster) is None: + return f"no {cluster}-CTA call counters for the {bucket.max_tokens}-token bucket" + return ( + f"operands outside the kernel: x {tuple(x.shape)} {x.dtype} contiguous={x.is_contiguous()}, " + f"residual contiguous={residual.is_contiguous()}, hidden {bucket.hidden}" + + ("" if x2 is None else f", x2 {tuple(x2.shape)} contiguous={x2.is_contiguous()}") + ) + + +def m3_mc_allreduce_norm( + x, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + x2=None, + quant=False, + prepushed=0, + drafter_norm=False, + carveout=True, + capture=None, +): + """Return (normed, updated_residual) or None when unsupported. ``carveout`` False drops opts bit 8 (the max-smem + carveout) for this call. ``drafter_norm``: the Eagle3 drafter boundary + (opts bit 16: normalize the unrounded FP32 sum, as the drafter's stock FlashInfer fused_add_rmsnorm does); 8-token + bucket only, x only, no MXFP8 copy. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture`` = (out [>= M, H] view, fc_norm weight + [H], eps): the boundary also writes the Eagle3 capture of the updated residual, bit for bit + m3_spec_tail.capture_norm (CAPTURE_NORM_IN_BOUNDARY; not with ``drafter_norm``). + + A None makes the caller take another collective, so it must be the same on every TP rank. It is: it depends only + on the shapes, dtypes and layouts of the operands (the same code builds them on every rank) and on whether the + workspace's flags exist, which ``get_workspace`` creates collectively, before any capture, for cluster 8.""" + workspace = _bucket(workspace, x.shape[0]) + if workspace is None: + return None + if ( + x.dim() != 2 + or x.shape[0] > MAX_TOKENS + or x.dtype != torch.bfloat16 + or x.shape[1] != workspace.hidden + or not x.is_contiguous() + or not residual.is_contiguous() + or (x2 is not None and (not x2.is_contiguous() or x2.shape != x.shape)) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + if drafter_norm and ( + workspace.max_tokens != 8 or x2 is not None or quant or prepushed not in (0, 2) + ): + return None + outs = _allreduce_norm_op( + x, + x2, + residual, + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and x.shape[1] % 128 == 0), + int(prepushed), + (_ar_opts(x.shape[0], workspace.max_tokens) & (~0 if carveout else ~8)) + | (16 if drafter_norm else 0), + *_capture_args(capture, drafter_norm), + own_op=own_op(), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def m3_mc_finalize_allreduce_norm( + fc2, + perm, + scales, + shared, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + quant=False, + capture=None, +): + """MoE finalize (sum_k scale * fc2[perm] + shared) + all-reduce + residual + RMSNorm in one kernel. + Returns (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture``: as ``m3_mc_allreduce_norm``.""" + m = residual.shape[0] + workspace = _bucket(workspace, m) + if workspace is None: + return None + if ( + residual.dim() != 2 + or m > MAX_TOKENS + or residual.shape[1] != workspace.hidden + or fc2.dtype != torch.bfloat16 + or not fc2.is_contiguous() + or fc2.dim() != 2 + or perm.dtype != torch.int32 + or perm.dim() != 2 + or perm.shape[0] != m + or not perm.is_contiguous() + or scales.dtype not in (torch.float32, torch.bfloat16) + or not scales.is_contiguous() + or (shared is not None and (shared.shape != residual.shape or not shared.is_contiguous())) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _finalize_allreduce_norm_op( + fc2, + perm, + scales, + shared, + residual.contiguous(), + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and residual.shape[1] % 128 == 0), + _ar_opts(m, workspace.max_tokens, finalize=True), + *_capture_args(capture, False), + own_op=own_op(), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +@functools.lru_cache(maxsize=None) +def _own_op_loaded() -> bool: + return getattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_mc_allreduce_norm", None) is not None + + +# The boundary all-reduces run this target's op (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm) rather than the +# image's: the same exchange and arithmetic built for the two-tray TP8 group. +OWN_OP = True +OWN_OP_PATH = "m3_mc_ar_own_op" + + +def own_op() -> bool: + """OWN_OP when this target's op is loaded (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm), with + its engagement line; else the image's op. The same library is loaded on every rank, so the ranks agree.""" + if not OWN_OP: + return False + if not _own_op_loaded(): + m3_report.missed( + OWN_OP_PATH, + None, + "trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm is not loaded (the image's op runs)", + ) + return False + m3_report.taken(OWN_OP_PATH, None, "the boundary all-reduces run this target's op") + return True + + +@functools.lru_cache(maxsize=None) +def capture_supported() -> bool: + """Whether the multicast all-reduce op the boundaries call (own_op) has the capture epilogue (cap_out / cap_w / + cap_eps). The same library is loaded on every rank, so the ranks agree.""" + name = ( + "minimax_m3_sm107_tp8_mc_allreduce_norm" + if own_op() + else "minimax_m3_sm107_tp2_mc_allreduce_norm" + ) + op = getattr(torch.ops.trtllm, name, None) + if op is None: + return False + return any(a.name == "cap_out" for a in op.default._schema.arguments) + + +def _capture_args(capture, drafter_norm: bool) -> tuple: + """(cap_out, cap_w, cap_eps) for the op, from ``capture`` = (out, fc_norm weight, eps) or None.""" + if capture is None: + return (None, None, 0.0) + assert not drafter_norm, "the capture epilogue is not built for the drafter boundary" + out, weight, eps = capture + return (out, weight, float(eps)) + + +_WORKSPACE: dict = {} +_AGREE_COMMS: dict = {} + + +def _tp_comm(mapping): + """A communicator over mapping's TP group, created once (collective over that group only).""" + key = tuple(mapping.tp_group) + comm = _AGREE_COMMS.get(key) + if comm is None: + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + comm = mapping.tp_group_pg + else: + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + _AGREE_COMMS[key] = comm + return comm + + +def tp_agree(mapping, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of mapping's TP group (one blocking host collective). + + For a fast path whose availability is decided from rank-local state (a kernel build that can fail on one rank): + every rank must take the same path, so that they issue the same collectives with the same exchange protocol and + the same kernels. Call it once per decision, at build time, at the same point on every rank, outside CUDA-graph + capture; never per step. With one rank, or attention DP (whose ranks run different batches), it returns ``ok``. + """ + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return ok + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("tp_agree must run outside CUDA-graph capture") + from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded + + return _mnnvl_workspace_all_succeeded(_tp_comm(mapping), bool(ok)) + + +def tp_agree_built(mapping, built, what: str): + """``built`` when it is not None on every TP rank, else None on every rank (see ``tp_agree``).""" + if tp_agree(mapping, built is not None) or built is None: + return built + from tensorrt_llm.logger import logger + + logger.warning(f"{what}: built here but not on another TP rank; every rank falls back") + return None + + +def get_workspace(mapping, hidden: int): + """Process-wide workspace (created collectively on first eager use); None while capturing before creation. + + The first eager call for a hidden size (the warmup forward, before any CUDA-graph capture) creates one + McAllReduceWorkspace per bucket in BUCKETS on every TP rank (MPI collective over the TP group) plus each + bucket's 8-CTA-cluster call counters; later calls, captured or not, return the same McBuckets. The workspaces + live until process exit. + """ + ws = _WORKSPACE.get(hidden) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + return None + wss = [McAllReduceWorkspace.create(mapping, hidden, b) for b in BUCKETS] + for w in wss: + w.flags_for(8) + ws = McBuckets(wss) + _WORKSPACE[hidden] = ws + return ws diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp8/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..d1730239463c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3841 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP rank: ``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 # capacity: fixed: the checkpoint's hidden size +# routed intermediate per rank: 768 at TP4, 384 at TP8 (FC2 K = 1.5 256-K tiles: the last tile's A box is half out +# of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 MMAs / scale atoms are issued; no padding, no extra +# bytes) +I_TP = _cfg("i_tp", 768) +TWO_I = 2 * I_TP +assert I_TP % 128 == 0 and 128 <= I_TP <= 768, I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 # capacity: fixed: the CTA size +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# EARLY_AR (scheduling only): when the finalize pushes into the multicast all-reduce, EVERY CTA releases the +# dependents right after the grid dependency, so the all-reduce (which reads only the pushed tagged words, without a +# grid wait) and, through its own early release, the next layer's QKV GEMV become resident on each SM this kernel +# frees, instead of after the last CTA's teardown. The arithmetic is unchanged. +EARLY_AR = bool(_cfg("early_ar", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +# the shared intermediate per rank: 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate positions: its +# gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP4: 2 x (32 + 32), the order +# of m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. +SH_I = _cfg("sh_i", 768) +SH_ROWS = SH_I // 24 +# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) +SH_SPLIT = 128 // (2 * SH_ROWS) +SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N +SH_KB = H // 32 # capacity: fixed: a kernel tile constant (gate_up K blocks) +SH_DKB = SH_I // 32 # down K blocks +SH_DKT = SH_I // 128 # down k-tiles +SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the +# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction +# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) +) # (in-chain: a win at M <= 2, not 4..8) +if SHARED: + assert SH_I in (384, 768), SH_I + assert SH_PRETILE or SH_SPLIT == 2, ( + "the un-pretiled shared stream is the 32-row (TP4) layout only" + ) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # capacity: fixed: pipeline stages (6 stages per unit) +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) +# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) +KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 + + +def _nkb2(t2: int) -> int: + """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" + return NUM_KB if t2 < KT2 - 1 else KB2_LAST + + +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy +# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few +# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. +HB_REPL = _cfg("hb_repl", 1) +assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) +HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # capacity: fixed: a warp role index +S_WARP = 6 # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +# E0_M4: at M = 4 the MMA warp routes +# token 3 (NRW0 = 4), so the early group-0 stages go to FC2 B builder warp 8 instead (idle until the routing is done; +# RW_EXTRA = 0 at M = 4). Issue order only: exact. +E0_M4 = ( + bool(_cfg("e0_m4", 1)) and NRW0 == 4 and RW_EXTRA == 0 and THREADS // 32 > 8 +) # M == 4 only: at M > 4 warp 8 routes (RW_EXTRA) +E0_WARP = ( + 8 if E0_M4 else MMA_WARP +) # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and (NRW0 <= 3 or E0_M4) and bool(_cfg("e0_off", 1)) +# B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing +# warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the +# two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. +if B1_AFTER_TOPK and XSTAGE and EARLY0 and not E0_OFFLOAD: + raise ValueError( + "m3_moe: b1_after_topk=1 with xstage, early0 and no early-stage offload deadlocks (warp 4 waits on x_bar[1] " + "before TOPK2_BAR, the B1 build waits on TOPK2_BAR before x_bar[1]); set b1_after_topk=0 or early0=0" + ) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 # capacity: fixed: the SM's tensor-memory columns +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +assert M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 +_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 +_SH_FIX = ( + SH_GT * SH_NB * 128 + + SH_DKT * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + _SH_XCH_BYTES + + 2 * SH_ROWS * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * SH_NB * 128 +_SB_XC = _SB_BH + SH_DKT * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + _SH_XCH_BYTES +_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = SH_NB +SH_SFA = SH_ACC_DN + 16 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +SH_TMEM = 256 # capacity: fixed: the shared CTAs' TMEM columns +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + # (a partial last k-tile has no atoms past I_TP) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SH_NB, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if cutlass.const_expr(SH_SPLIT == 2): + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + else: + # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are + # accumulator columns 8 qd + t + av = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_GU + qd * cutlass.Int32(8), + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(1): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd == cutlass.Int32(0): + # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + for s_ in cutlass.range_constexpr(1, SH_SPLIT): + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + lane * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(SH_ROWS * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(SH_DKT): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(EARLY_AR): + if ar_on: + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + for rr in cutlass.range_constexpr(HB_REPL): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4) + + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + hb_off = cutlass.Int32(0) + if cutlass.const_expr(HB_REPL > 1): + # this CTA's copy of the records + hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview( + (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) + ), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if cutlass.const_expr(HB_REPL == 1): + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + else: + if lane < HB_REPL: # copy `lane` + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4) + + lane * cutlass.Int32(HB_STRIDE), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/op.py new file mode 100644 index 000000000000..70deb357f4a8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/op.py @@ -0,0 +1,693 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP8 rank (routed and shared intermediate 384 per rank), +over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +INTER = 384 # capacity: fixed: the routed experts' width per tp8 rank +NUM_EXPERTS = 128 +SHARED_INTER = 384 # capacity: fixed: the shared expert's width per tp8 rank (the routed one) +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size (m3_moe's shared-expert CTAs (beside num_ctas routed CTAs)) +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +U1 = 2 * INTER // 32 # FC1 units per expert +# The router GEMV takes at most 2 experts per CTA (router weights held in registers). +MIN_CTAS = 64 # capacity: fixed: an SM split, not a size + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # capacity: fixed: pipeline stages (FC1 stages (items) per unit) +_ITEM_BYTES = _UNIT_BYTES // _S1 # capacity: fixed: a ring item's bytes +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters (7 and 8 tokens). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 1, 1) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit). 6 at 8 tokens measured 0.4-0.5 us faster than the +# default 7 (same MMAs, bit-identical outputs). +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 6) +# ``hb_repl``: copies of the tagged FC1 -> FC2 records (applied when > 1); CTA bx polls copy bx % hb_repl, which +# spreads the polls of the few hot record lines over L2. 7 and 8 tokens hand off through the B images instead. +HB_REPL_BY_M: Tuple[int, ...] = (32, 32, 8, 8, 1, 1, 1, 1) +# ``lg_repl``: replicas of the router logit exchange at up to LG_REPL_MAX_TOKENS tokens (fewer stores before the +# exchange; bit-identical outputs). Above that the kernel's default applies. +LG_REPL = 32 +LG_REPL_MAX_TOKENS = 4 # capacity: fixed: measured crossover +# Byte-balanced FC1 ranges below one unit per CTA (:func:`_fine_ranges`) for groups too small to give every CTA a +# whole unit; False gives such groups whole units on the CTAs with the least FC2 work. +FINE_PLAN = True +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 + = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row + R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then + down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per + layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + P = I_ // n_cta + split = 128 // (2 * P) + gt = H_ // 128 // split + dkt = I_ // 128 + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + octant = r // P + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(2 * dkt): + jt, kt = gd // dkt, gd % dkt + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, gt + gd, :16384] = img.reshape(-1) + out[s, gt + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def _fine_ranges(weights, total: int): + """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item + counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary + strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit + runs at least to that unit's end. Returns the item counts.""" + n = len(weights) + items = [0] * n + s = 0 + acc = 0.0 + for c in range(n): + acc += weights[c] + e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) + if s % _S1 and e < total: # started inside a unit: run to (at least) its end + e = max(e, min(total, (s // _S1 + 1) * _S1)) + items[c] = e - s + s = e + assert s == total + return items + + +def make_phase_plan( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = FINE_PLAN, +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [max(_S1, (limit - f) // _ITEM_BYTES) for f in fc2] + + lo, hi = 0, max(fc2) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 + lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] + sw = sum(wts) + items = _fine_ranges([w * total / sw for w in wts], total) + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + inter: int = INTER, + early_release: bool = True, + ): + """``inter``: the routed intermediate per rank; this engine is built for INTER (384) only. + ``early_release``: on a call that pushes into the boundary all-reduce (``arp``), every CTA releases the PDL + dependents right after its grid dependency wait (the kernel's EARLY_AR); False releases them after the push.""" + if inter != INTER: + raise ValueError( + f"m3_moe (TP8): routed intermediate {inter} per rank is not supported (only {INTER})" + ) + self.inter = inter + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, num_tokens, M3_MOE_LAG, 0.0 if num_tokens == 1 else 1.0, True + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "i_tp": INTER, + } + if HB_REPL_BY_M[num_tokens - 1] > 1: + self.config["hb_repl"] = HB_REPL_BY_M[num_tokens - 1] + if num_tokens <= LG_REPL_MAX_TOKENS: + self.config["lg_repl"] = LG_REPL + if shared: + self.config["sh_i"] = SHARED_INTER + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if not early_release: + self.config["early_ar"] = 0 + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + # E0@m4: at 4 tokens the early group-0 FC1 stages are issued by FC2 B builder warp 8 instead of the routing + # warp 4 (issue order only, bitwise; TP4 -0.19 / TP8 -0.21 us/call). + if num_tokens == 4: + from tensorrt_llm.logger import logger as _e0log + + _e0log.info( + "MiniMax-M3 MoE engine E0M4 (%s): early group-0 stages on warp 8, tokens=4" % "TP8" + ) + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full( + (mod.G_CAP * mod.U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw + ) + # the FC1 -> FC2 records; with HB_REPL > 1 copies, HB_STRIDE int32 words apart + hb_words = mod.G_CAP * num_tokens * mod.U1 * 4 + if mod.HB_REPL > 1: + hb_words = mod.HB_REPL * mod.HB_STRIDE + self.hbuf32 = torch.zeros(hb_words, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [768, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 384], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + key = (_ENGINE_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp8/kernels/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_report.py new file mode 100644 index 000000000000..7f864bd240c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_report.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast-path reports a kernel module makes from inside a call: ``taken`` (an engagement line), ``missed`` (a covered +shape that cannot run: raises under ``modeling_v2='require'``), ``uncovered`` and ``capture_missed``. + +The policy (logging once, raising under require) is the target's: ``modeling.py`` installs its functions here at import, +so a kernel reports without importing the target module. Until then every report is a no-op. +""" + +from __future__ import annotations + +from typing import Callable, Optional + + +def _noop(*args, **kwargs) -> None: + return None + + +taken: Callable[..., None] = _noop +missed: Callable[..., None] = _noop +uncovered: Callable[..., None] = _noop +capture_missed: Callable[..., None] = _noop + + +def install( + *, + on_taken: Callable[[str, Optional[int], str], None], + on_missed: Callable[[str, Optional[int], str], None], + on_uncovered: Callable[[str, str], None], + on_capture_missed: Callable[[str, Optional[int], str], None], +) -> None: + """Route this module's reports to the target's policy.""" + global taken, missed, uncovered, capture_missed + taken, missed, uncovered, capture_missed = on_taken, on_missed, on_uncovered, on_capture_missed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_router_gemv.py new file mode 100644 index 000000000000..30855eb7f05d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/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 # capacity: fixed: the device's shared-memory limit + + +def m3_router_gemv( + x: torch.Tensor, weight: torch.Tensor, rows: int = 1, use_pdl: bool = False +) -> Optional[torch.Tensor]: + """fp32 ``x @ weight.T`` for bf16 ``x`` with at most 16 rows; None when unsupported (only ``rows=1`` is built).""" + if ( + rows != 1 + or not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or weight.shape[0] % rows != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or rows * x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return _router_gemv_op(x, weight, rows, use_pdl) + + +def m3_router_gemv_fp4( + x: torch.Tensor, weight: torch.Tensor, sf_scale: torch.Tensor, use_pdl: bool = False +) -> Optional[tuple]: + """Router logits plus the NVFP4 activation the routed experts consume (bit-identical to + ``fp4_quantize(x, sf_scale, 16, False, False)``): (logits [M, E] fp32, x_fp4 [M, K / 2] uint8, + x_sf [M * K / 16] uint8 linear). None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or sf_scale is None + or sf_scale.dtype != torch.float32 + or sf_scale.numel() != 1 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return tuple(_router_gemv_fp4_op(x, weight, sf_scale.reshape(1), use_pdl)) + + +def m3_router_tc( + x: torch.Tensor, + weight: torch.Tensor, + sf_scale: Optional[torch.Tensor] = None, + use_pdl: bool = False, +) -> Optional[tuple]: + """Router logits (TF32 tensor cores, like the native TF32 GEMM) for 1..48 tokens with 128 experts, plus the + NVFP4 activation when ``sf_scale`` is given. Returns (logits,) or (logits, x_fp4, x_sf); None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 48 + or weight.dtype != torch.float32 + or weight.shape[0] != 128 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or (x.stride(0) * 2) % 16 != 0 + or x.data_ptr() % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or 16 * (x.shape[1] // 8 + 4) * 4 + x.shape[0] * (x.shape[1] // 8 + 8) * 2 + > _MAX_DYNAMIC_SMEM + or (sf_scale is not None and (sf_scale.dtype != torch.float32 or sf_scale.numel() != 1)) + ): + return None + res = _router_tc_op(x, weight, None if sf_scale is None else sf_scale.reshape(1), use_pdl) + return tuple(res) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_sampled_accept.py new file mode 100644 index 000000000000..c3c2edd5d078 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_sampled_accept.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). + +Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +active must stay on the generic path; the caller decides. + +Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +(the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens heavier +than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the nucleus +with probability >= top_p, so most rows decide on their second exchange. + +Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based Philox +keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one 4-word +record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and reads +the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the same +order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing and +two call halves keep a rank one call ahead from overwriting a slot still being read. +""" + +from typing import Optional + +import torch + +from . import m3_report + +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu) and called as +# torch.ops.trtllm.minimax_m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) + + +def _op_args() -> frozenset: + """The optional arguments the loaded trtllm::minimax_m3_sm107tp8_sampled_accept takes (the tp2 op is used as it + is, read-only: an older image's op has neither prewait_params nor early_trigger).""" + global _OP_ARGS + if _OP_ARGS is None: + schema = torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept.default._schema + _OP_ARGS = frozenset(a.name for a in schema.arguments) + return _OP_ARGS + + +_OP_ARGS = None + + +def one_round_k() -> int: + """The SA-1R variant (ONE_ROUND_K) when the loaded op takes it, else 0 (two rounds).""" + return int(ONE_ROUND_K) if "one_round_k" in _op_args() else 0 + + +# The vocab shard limit of trtllm::minimax_m3_sm107tp8_sampled_accept (kMaxVs, kept from the TP2 kernel it copies): +# 200064 / 2 = 100032 columns. +MAX_VS = 100352 # capacity: fixed: the vocabulary shard bound + + +class SampledAcceptWorkspace: + """Exchange slots behind one multicast mapping (two call halves), the (finished rows, call) counters, the per-row + token scratch, an extra-attempt counter and the SA-1R round-0 decided counter. Created collectively, outside + CUDA-graph capture.""" + + def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): + self.uc, self.mc_ptr, self.rank, self.world = uc, mc_ptr, rank, world + self.handle, self.comm = handle, comm + dev = uc.device + self.counters = torch.zeros(2, dtype=torch.int32, device=dev) + self.tokens = torch.zeros(MAX_ROWS, dtype=torch.int32, device=dev) + self.extra_attempts = torch.zeros(1, dtype=torch.int32, device=dev) + self.decided = torch.zeros(1, dtype=torch.int32, device=dev) + + @classmethod + def create(cls, mapping, device) -> "SampledAcceptWorkspace": + # SA-1R's records carry the round-0 list words: the size depends on the variant + words = int( + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words( + max(mapping.tp_size, 1), one_round_k() + ) + if one_round_k() + else torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words( + max(mapping.tp_size, 1) + ) + ) + if mapping.tp_size == 1: + uc = torch.zeros(words, dtype=torch.int64, device=device) + return cls(uc, uc.data_ptr(), 0, 1) + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 8, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("SampledAcceptWorkspace.create failed on another rank") + return cls(uc, mc.data_ptr(), mapping.tp_rank, mapping.tp_size, handle, comm) + + +_WS: Optional[SampledAcceptWorkspace] = None + + +def workspace(mapping, device) -> Optional[SampledAcceptWorkspace]: + global _WS + if _WS is None: + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("spec_tail", None, "the sampled-accept workspace") + return None + _WS = SampledAcceptWorkspace.create(mapping, device) + return _WS + + +def sampled_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: Optional[torch.Tensor], + seeds: torch.Tensor, + offsets: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, + ws: Optional[SampledAcceptWorkspace] = None, + use_pdl: bool = True, +): + """Sampled verification (temperature + top-p) from this rank's bf16 logits shard [nreq * (draft_len + 1), V / tp]: + outputs as shard_accept (accepted [nreq, D+1], num_accepted [nreq], gather_ids [nreq]); None + when unavailable. Per-row sampling parameters and Philox (seed, offset) are the spec metadata's, one per row. + top_ks / min_ps are not applied: rows with top-k or min-p active must be routed to the generic path by the + caller.""" + del top_ks, min_ps + nreq, draft_len = draft_tokens.shape + rows = nreq * (draft_len + 1) + if logits.dim() != 2 or logits.shape[0] != rows or rows > MAX_ROWS or logits.shape[1] > MAX_VS: + return None + if logits.dtype != torch.bfloat16 or logits.stride(1) != 1: + return None + ws = ws if ws is not None else workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept( + logits, + int(vocab_offset), + temperatures[:rows].float().contiguous(), + top_ps[:rows].float().contiguous(), + seeds[:rows].long().contiguous(), + offsets[:rows].long().contiguous(), + ws.uc, + ws.mc_ptr, + ws.counters, + ws.tokens, + ws.rank, + ws.world, + draft_tokens.int().contiguous(), + batch_idx.int().contiguous(), + accepted, + num_accepted, + gather_ids, + int(force_total), + next_new, + force_vec, + ws.extra_attempts, + bool(use_pdl), + **{ + name: bool(on) + for name, on in ( + ("prewait_params", PREWAIT_PARAMS), + ("early_trigger", EARLY_TRIGGER), + ) + if name in _op_args() + }, + **({"one_round_k": one_round_k(), "decided": ws.decided} if one_round_k() else {}), + **({"ilp_walk": True} if _ilp_walk(rows) else {}), + ) + if one_round_k(): + m3_report.taken("m3_sampled_accept_1r", rows, f"one_round_k {one_round_k()}") + return accepted, num_accepted, gather_ids + + +# Off at TP8: the sampled verify loads each row's sampling parameters after its PDL dependency wait, and releases its +# dependents at the acceptance epilogue (both are timing only; the same values and tokens either way). +PREWAIT_PARAMS = False +EARLY_TRIGGER = False +# SA-1R: each CTA's round-0 record also carries its largest logits (82: each thread's top 2 pooled, the CTA's top 8 +# plus the bound of everything it does not list), so a row whose first candidate is provably in the nucleus decides in +# round 0. The same tokens, accepted counts and Philox draws as the two-round kernel, bit for bit; 0 = two rounds. +ONE_ROUND_K = 82 # capacity: fixed: a record format, not a size +# SA-PW: the owner walk takes its tokens 8 at a time (the same draw, tokens and Philox draws as the one-token walk, +# bit for bit). Off. +ILP_WALK = False +ILP_WALK_PATH = "m3_sampled_accept_ilp_walk" + + +def _ilp_walk(rows: int) -> bool: + """ILP_WALK when the loaded op takes ilp_walk, logged once per row count as taken.""" + if not ILP_WALK or "ilp_walk" not in _op_args(): + return False + m3_report.taken( + ILP_WALK_PATH, + rows, + "the sampled verify's owner walk takes 8 tokens at a time", + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_select_long_dsl.py new file mode 100644 index 000000000000..10bd966adf98 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_select_long_dsl.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). + +Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long``, +one 1024-thread CTA per row), from a CTA of 384 threads: + +* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and + mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's + valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned + key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped + block and masks the result, so a thread's loads issue back to back. +* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of + the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. +* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by + that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate + against the whole row instead. +* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. + +The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait +(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are +compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and +loads them before the executor's warmup, so a decode step never compiles one. +""" + +from typing import Dict, Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.utils import ( + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +# The 384-thread select reads its row's valid-block count before its grid wait instead of after it, which takes one +# dependent global load off the post-wait path; the selection is the same either way. Exact: the count +# (MsaMetadata.msa_n_valid_blocks) is written only before the step's first layer, behind a plain launch. +SELECT_N_VALID_PREWAIT = True + +SELECT_THREADS = 384 +SELECT_WARPS = SELECT_THREADS // 32 +# 1M tokens + draft tokens at 128 tokens per block. +MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row +MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging +TOPK = 16 # capacity: fixed: MSA's top-k +NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) +# The key of -inf (the reference's 0xff800000): never emitted. +NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 +# (keys per thread, rows up to this many valid blocks) +ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) + +# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare +# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group +# minima [16], per-warp candidate counts [SELECT_WARPS]. +_CAND = 0 +_PICK_KEY = _CAND + 2 * MAX_CAND +_PICK_BLOCK = _PICK_KEY + TOPK +_WARP_MIN = _PICK_BLOCK + TOPK +_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 +_WARP_CNT = _GROUP_MIN + 16 +SMEM_WORDS = _WARP_CNT + SELECT_WARPS +SMEM_BYTES = 4 * SMEM_WORDS + + +def _key(score): + """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" + bits = cutlass.Float32(score).bitcast(cutlass.Int32) + return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) + + +def _imin(a, b): + return cutlass.select_(b < a, b, a) + + +def _imax(a, b): + return cutlass.select_(b > a, b, a) + + +def _forced_key(scores, head, block, query, init_blocks, local_start): + raw = scores[head, block, query] + s = cutlass.select_( + block >= local_start, + cutlass.Float32(LOCAL_SCORE), + cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), + ) + return _key(s) + + +@cute.jit +def _select_row( + row, + scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides + n_valid: cute.Tensor, # int32 [>= total_q] + out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] + init_blocks, + local_blocks, + num_kv_heads, + smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory + bar_id: cutlass.Constexpr, + loaded_trigger: cutlass.Constexpr, + raw_prewait, + n_valid_prewait: cutlass.Constexpr, +): + """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier + ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the + row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by + the caller before its grid wait; otherwise it is read here.""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % cutlass.Int32(32) + warp = tid // cutlass.Int32(32) + query = row // num_kv_heads + head = row % num_kv_heads + if cutlass.const_expr(n_valid_prewait): + raw = raw_prewait + else: + raw = n_valid[query] + valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) + local_start = _imax(raw - local_blocks, cutlass.Int32(0)) + cand = cute.make_tensor( + cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) + ) + + if tid < TOPK: + smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) + smem[_PICK_BLOCK + tid] = cutlass.Int32(0) + + if valid <= TOPK: + # Every valid block is picked (thread t holds block t). + if tid < valid: + smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) + smem[_PICK_BLOCK + tid] = tid + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): + items = ITEM_CLASSES[ci][0] + lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] + hi = ITEM_CLASSES[ci][1] + if (valid > lo) & (valid <= hi): + key = cute.make_rmem_tensor(items, cutlass.Int32) + local_min = cutlass.Int32(NO_KEY) + for s in cutlass.range_constexpr(items): + # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. + block = tid + cutlass.Int32(s * SELECT_THREADS) + k = _forced_key( + scores, head, _imin(block, valid - 1), query, init_blocks, local_start + ) + key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) + local_min = _imin(local_min, key[s]) + local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) + if lane < 16: + smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + if tid < 16: + g = cutlass.Int32(NO_KEY) + for w in cutlass.range_constexpr(SELECT_WARPS): + g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) + smem[_GROUP_MIN + tid] = g + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + bound = cutlass.Int32( + -2147483648 + ) # the smallest key: signed keys, unlike the reference's unsigned ones + for g in cutlass.range_constexpr(16): + bound = _imax(bound, smem[_GROUP_MIN + g]) + + # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(items): + cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) + incl = cnt + for d in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync( + incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) + ) + incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) + if lane == 31: + smem[_WARP_CNT + warp] = incl + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + before = cutlass.Int32(0) + total = cutlass.Int32(0) + for w in cutlass.range_constexpr(SELECT_WARPS): + c = smem[_WARP_CNT + w] + before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) + total += c + slot = before + incl - cnt + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + if slot < MAX_CAND: + cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( + tid + cutlass.Int32(s * SELECT_THREADS) + ) + slot += 1 + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if total <= MAX_CAND: + for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): + mine = cand[i] + rank = cutlass.Int32(0) + for j in cutlass.range(total, unroll=8): + rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) + smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) + else: + # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + block = tid + cutlass.Int32(s * SELECT_THREADS) + rank = cutlass.Int32(0) + for j in cutlass.range(valid, unroll=1): + kj = _forced_key(scores, head, j, query, init_blocks, local_start) + less = (kj < key[s]) | ((kj == key[s]) & (j < block)) + rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = key[s] + smem[_PICK_BLOCK + rank] = block + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if warp == 0: + # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. + v = cutlass.Int32(NO_KEY) + if lane < TOPK: + pk = smem[_PICK_KEY + lane] + if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): + v = smem[_PICK_BLOCK + lane] + for size_log in cutlass.range_constexpr(1, 6): + for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): + other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) + ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) + lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) + v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) + if lane < TOPK: + out[query, head, lane] = cutlass.select_( + v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v + ) + + +class _SelectLongDslKernel: + """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count + before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because + ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's + first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it + (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" + + def __init__( + self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool + ): + self.use_pdl = use_pdl + self.early_trigger = early_trigger + self.loaded_trigger = (not early_trigger) and trigger_phase == 1 + self.n_valid_prewait = n_valid_prewait + + @cute.kernel + def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): + row, _, _ = cute.arch.block_idx() + smem = SmemAllocator() + words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) + raw_prewait = cutlass.Int32(0) + if cutlass.const_expr(self.n_valid_prewait): + raw_prewait = n_valid[row // num_kv_heads] + griddepcontrol_wait() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + _select_row( + row, + scores, + n_valid, + out, + init_blocks, + local_blocks, + num_kv_heads, + words, + bar_id=1, + loaded_trigger=self.loaded_trigger, + raw_prewait=raw_prewait, + n_valid_prewait=self.n_valid_prewait, + ) + + @cute.jit + def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): + num_kv_heads = scores.shape[0] + self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( + grid=(scores.shape[2] * num_kv_heads, 1, 1), + block=(SELECT_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + +#: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). +kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} + + +def kernel_key( + use_pdl: bool, early_trigger: bool, trigger_phase: int +) -> Tuple[bool, bool, int, bool]: + """The compile key of a launch (the trigger phase does not matter with the early trigger), with this target's + ``SELECT_N_VALID_PREWAIT``.""" + return ( + bool(use_pdl), + bool(early_trigger), + 0 if early_trigger else int(trigger_phase), + bool(SELECT_N_VALID_PREWAIT), + ) + + +def _compile(key: Tuple[bool, bool, int, bool]): + if key not in kernel_cache: + heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() + scores = make_fake_tensor( + cutlass.Float32, + (heads, blocks, rows), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) + out = make_fake_tensor( + cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) + ) + kernel_cache[key] = cute.compile( + _SelectLongDslKernel(*key), + scores, + n_valid, + out, + cutlass.Int32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + return kernel_cache[key] + + +def select_long_dsl( + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool, + early_trigger: bool, + trigger_phase: int, +) -> None: + """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" + _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( + scores, n_valid_blocks, out, init_blocks, local_blocks + ) + + +def prebuild(keys, device) -> list: + """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" + scores = torch.zeros( + 1, 2048, 4, dtype=torch.float32, device=device + ) # capacity: fixed: a dummy prebuild row, not a size + n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) + out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) + built = [] + for key in keys: + key = kernel_key(*key) + _compile(key)(scores, n_valid, out, 0, 1) + built.append(key) + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_shard_accept.py new file mode 100644 index 000000000000..8caaee389d0a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/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 m3_report + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 # capacity: fixed: the kernel's built verify rows + + +@dataclass(eq=False) +class AcceptWorkspace: + """Per-rank exchange slots ([2 halves][SHARD_ROWS rows][world] uint64) behind one multicast mapping, the local + key accumulators and the (finished CTAs, call) counters.""" + + uc: torch.Tensor + mc_ptr: int + keys: torch.Tensor + counters: torch.Tensor + rank: int + world_size: int + handle: Any = None + comm: Any = None + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + if not _mnnvl_workspace_all_succeeded(comm, True): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = mc.data_ptr(), mapping.tp_rank + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None + return AcceptWorkspace( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) + + +def _workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace, created collectively outside CUDA-graph capture on first use.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("m3_shard_accept", None, "the shard-accept workspace") + return None + _WS = create_workspace(mapping, device, SHARD_ROWS) + return _WS + + +def shard_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, +): + """Greedy verification from this rank's BF16 logits shard. + + Args: + logits: [nreq * (draft_len + 1) <= SHARD_ROWS, V / tp] BF16, unit column stride: this rank's columns + (global ids ``vocab_offset ..``) of the verify rows, request-major. + draft_tokens: [nreq, draft_len] int32, the drafts being verified. + batch_idx: [nreq] int32, each request's row in the verify layout. + force_total / force_vec: forced acceptance (see ``m3_shard_accept``). + next_new: optional int32 [nreq, >= 1]; column 0 receives the last accepted token. + + Returns: + (accepted_tokens [nreq, draft_len + 1] int32, num_accepted [nreq] int32, first-pass gather ids [nreq] + int64), or None when the shape is not served or the workspace is missing during a capture. Collective over + the TP group. + """ + nreq, draft_len = draft_tokens.shape + if ( + logits.dim() != 2 + or logits.shape[0] != nreq * (draft_len + 1) + or logits.shape[0] > SHARD_ROWS + or logits.dtype != torch.bfloat16 + or logits.stride(1) != 1 + ): + return None + ws = _workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_tokens.contiguous(), + batch_idx, + accepted, + num_accepted, + gather_ids, + force_total, + next_new, + _num_sms(), + force_vec, + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp8/kernels/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/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_107__tp8/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_107__tp8/kernels/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_tokmap_guard.py new file mode 100644 index 000000000000..db0077117fe3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_tokmap_guard.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The token map of ``thop.trtllm_gen_generation_preprocess`` for one request with several tokens. + +For multi-token generation the op's QKV preprocessing places every row by the workspace token map (``tokensInfo``, +(request, token) per row), but builds that map only when ``isBuildDecoderInfoKernelNeeded()`` holds, which at one +request with the default RoPE it does not. The op then applies whatever map the last multi-request call left in the +workspace: a row can land on another request's length and pages, and its RoPE position and K/V slot are wrong. This +module wraps the op so that, for one request and several tokens, the map the build would write, (0, t) for each row t, +is written into the op's own map slice of the workspace (the offset from the op's layout function) just before the +call. Every other call passes through unchanged. The rows' positions still come from the request's generation length, +exactly as with a freshly built map. + +The write is a device copy from a map prebuilt at install (no allocation, no host sync), so it is captured into the +decode graphs with the call it precedes. +""" + +from __future__ import annotations + +from typing import Callable, Dict, Optional, Tuple + +import torch + +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.logger import logger + +#: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). +_ARG_NAMES = ( + "qkv_input", + "workspace", + "sequence_lengths", + "spec_decoding_generation_lengths", + "spec_decoding_position_offsets", + "kv_cache_block_offsets", + "host_kv_cache_pool_pointers", + "host_kv_cache_pool_mapping", + "kv_scale_orig_quant", + "kv_scale_quant_orig", + "attention_output_orig_quant", + "rotary_inv_freq", + "rotary_cos_sin", + "mrope_position_deltas", + "layer_idx", + "seq_offset", + "num_heads", + "num_kv_heads", + "head_size", + "tokens_per_block", + "kv_cache_quant_mode", + "max_attention_window_size", + "cyclic_attention_window_size", + "num_tokens", + "batch_beam", + "input_seq_length", + "max_past_kv_length", + "rotary_embedding_dim", + "rotary_embedding_base", + "rotary_embedding_scale_type", + "rotary_embedding_scale", + "rotary_embedding_max_positions", + "position_embedding_type", + "bmm1_scale", + "bmm2_scale", + "fp8_context_fmha", + "predicted_tokens_per_seq", + "attention_chunk_size", + "multi_processor_count", + "total_num_blocks", + "kv_factor", + "need_build_kv_cache_metadata", + "cross_attention", + "skip_fmha_workspace", +) +_INT2_BYTES = 8 # capacity: fixed: an element size + +_original: Optional[Callable] = None +_max_rows = 0 +_maps: Dict[torch.device, torch.Tensor] = {} +_offsets: Dict[Tuple, int] = {} + + +def _no_taken(rows: int) -> None: + pass + + +def _no_missed(rows: int, reason: str) -> None: + pass + + +_on_taken: Callable[[int], None] = _no_taken +_on_missed: Callable[[int, str], None] = _no_missed + + +def _indexed(device: torch.device) -> torch.device: + """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: + """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" + device = _indexed(device) + m = _maps.get(device) + if m is None or m.shape[0] < rows: + reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" + _on_missed(rows, reason) + raise RuntimeError(f"m3_tokmap_guard: {reason}") + return m[:rows] + + +def _map_offset(args: dict, rows: int) -> int: + x = args["qkv_input"] + skip = bool(args.get("skip_fmha_workspace", False)) + key = ( + x.dtype, + rows, + int(args["num_heads"]), + int(args["head_size"]), + int(args["rotary_embedding_dim"]), + int(args["num_kv_heads"]), + skip, + ) + off = _offsets.get(key) + if off is None: + layout = thop.get_trtllm_gen_generation_workspace_layout( + x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip + ) + off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) + if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: + raise RuntimeError( + f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " + f"(offset {off}, {size} bytes): the layout changed under this guard" + ) + _offsets[key] = off + return off + + +def _guarded(*args, **kwargs): + a = dict(zip(_ARG_NAMES, args)) + a.update(kwargs) + rows = int(a["num_tokens"]) + if ( + a.get("spec_decoding_generation_lengths") is not None + and int(a["predicted_tokens_per_seq"]) > 1 + and int(a["batch_beam"]) == 1 + and rows > 1 + and not bool(a.get("cross_attention", False)) + ): + ws = a["workspace"] + src = _one_request_map(ws.device, rows) + off = _map_offset(a, rows) + # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. + dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] + dst.view(torch.int32).view(rows, 2).copy_(src) + _on_taken(rows) + return _original(*args, **kwargs) + + +def install( + device: torch.device, + max_rows: int, + on_taken: Callable[[int], None], + on_missed: Callable[[int, str], None] = _no_missed, +) -> bool: + """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to + ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this + call installed the wrapper.""" + global _original, _max_rows, _on_taken, _on_missed + device = _indexed(device) + if device not in _maps or _maps[device].shape[0] < max_rows: + t = torch.arange(max_rows, dtype=torch.int32, device=device) + _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() + _max_rows = max(_max_rows, max_rows) + _on_taken = on_taken + _on_missed = on_missed + if _original is not None: + return False + _original = thop.trtllm_gen_generation_preprocess + thop.trtllm_gen_generation_preprocess = _guarded + logger.info( + "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " + f"(up to {max_rows} rows)" + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp8/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/__init__.py new file mode 100644 index 000000000000..e6c3d806a4a3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/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 TP2 +projection shapes on SM107 (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..3ba6fb0d7c8c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1472 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +``prod`` (#117, the sparse QKV + index projection at split 8, at most 8 tokens): instead of storing the packed row, the +owners stage their bf16 rows, the CTAs exchange them by st.async so that rank t holds token t's 128-dim row of this +tile's packed head, and rank t's warp 4 runs that head's producer (``m3_layer.phase_producer.producer_head``: RMSNorm, +RoPE, FP8 q / index-q / index-K and NVFP4 K / V stores) on it, bit for bit what the separate producer kernel does with +the stored row. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +from ..m3_layer.phase_producer import ProducerArgs, bf16_bits_to_f32, producer_head + +CTA_M = 128 # capacity: fixed: a kernel tile constant (weight rows per CTA = the MMA's M) +MMA_N = 8 # token columns +CTA_K = 128 # capacity: fixed: a kernel tile constant (one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row) +MMA_INST_K = 32 +# 4 MMAs per k-tile, one per scale byte of the atom. +K_BLOCKS = CTA_K // MMA_INST_K # capacity: fixed: a kernel tile constant +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _ld_shared_v4_b32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.b32 of four int32 words at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.i32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", "=r,=r,=r,=r,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [cutlass.Int32(_llvm.extractvalue(_T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +# The swigluoai activation of MiniMax-M3's dense MLP, gate * sigmoid(1.702 * gate) * (up + 1) with gate clamped to <= 7 +# and up to [-7, 7], in exactly the instruction sequence the Triton silu_and_mul kernel compiles to (modules/swiglu.py +# on SM107): the gate clamp on the packed bf16 pair, the up clamp as min.xorsign.abs, two rounded multiplies into +# ex2.approx, div.full for the reciprocal, two rounded multiplies, round to bf16. No multiply feeds an add, so the +# explicit .rn forms compile to what the Triton kernel's contractable forms do. +SWIGLU_OAI_ALPHA = 1.702 +SWIGLU_OAI_LIMIT = 7.0 + + +@dsl_user_op +def _swiglu_oai_bf16x2(gate_word, up_word, *, loc=None, ip=None): + """The activation of one packed bf16 pair of gate and up values: its two bf16 results as one int32 word (the low + half from the low halves).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [cutlass.Int32(w).ir_value(loc=loc, ip=ip) for w in (gate_word, up_word)], + "{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, $1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, $2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 $0, r1, r0;\n\t}", "=r,r,r", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[ + bool + ], # x_bf16: x_words is [M, 2K] gate | up; the activation is swigluoai(gate, up) + # prod: the producer's addresses and strides (ProducerArgs fields; `packed` is unused), then num_pages and eps. + pr_q: cutlass.Int64, + pr_iq: cutlass.Int64, + pr_ik: cutlass.Int64, + pr_kvd: cutlass.Int64, + pr_kvs: cutlass.Int64, + pr_loc: cutlass.Int64, + pr_quant: cutlass.Int64, + pr_pos: cutlass.Int64, + pr_rope: cutlass.Int64, + pr_qw: cutlass.Int64, + pr_kw: cutlass.Int64, + pr_iqw: cutlass.Int64, + pr_ikw: cutlass.Int64, + pr_ips: cutlass.Int64, + pr_its: cutlass.Int64, + pr_dps: cutlass.Int64, + pr_dls: cutlass.Int64, + pr_dts: cutlass.Int64, + pr_sps: cutlass.Int64, + pr_sls: cutlass.Int64, + pr_pages: cutlass.Int32, + pr_eps: cutlass.Float32, + prod: cutlass.Constexpr[bool], +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # prod: this CTA's bf16 owner rows by token ([8][rows_per_rank]), and the packed-head row of token `rank` gathered + # from the cluster (128 bf16 as 64 words, completed on pr_full by 8 x rows_per_rank * 2 bytes of st.async). + pr_stage = None + pr_row = None + pr_full = None + if cutlass.const_expr(prod): + pr_stage = cutlass.Array( + cutlass.Uint16, 8 * rows_per_rank, space=cutlass.AddressSpace.smem, alignment=16 + ) + pr_row = cutlass.Array(cutlass.Int32, 64, space=cutlass.AddressSpace.smem, alignment=16) + pr_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if cutlass.const_expr(prod): + # Completed by the 8 ranks' st.async (this one's included) only on a rank that receives a token row. + prims.mbarrier_init(pr_full, 1) + if rank < num_tokens: + prims.mbarrier_arrive_expect_tx(pr_full, split * rows_per_rank * 2) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in if x_swiglu else k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + if cutlass.const_expr(x_swiglu): + # The up half of the same columns, K values after the gate in the row. + qz_wvecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(k_in // 2 + qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + if cutlass.const_expr(x_swiglu): + qz_word = _swiglu_oai_bf16x2( + qz_word, cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + ) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The flags are rotated by the previous all-reduce after its dependent release; this load is ordered by + # the wait above, which waits for ALL in-flight upstream grids, not only the immediate predecessor (PTX + # griddepcontrol; MoE test_pdl_transitive_wait.py). Keep the wait before this load. + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.const_expr(prod): + # The bf16 the packed row would hold (RN, as the y store below), staged by token. + if cutlass.Int32(4 * g + j) < num_tokens: + pr_stage.store( + cutlass.Uint16( + (_bf16_bits(sums[j]) >> cutlass.Int32(16)) + & cutlass.Int32(0xFFFF) + ), + idx=cutlass.Int32((4 * g + j) * rows_per_rank) + local, + ) + else: + if cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The flags are rotated by the previous all-reduce after its dependent release; this load is ordered + # by the wait above, which waits for ALL in-flight upstream grids, not only the immediate + # predecessor (PTX griddepcontrol; MoE test_pdl_transitive_wait.py). Keep the wait before this load. + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(prod): + # Transpose: thread e < 2 T sends half e % 2 (16 bytes = 8 dims) of token e / 2's staged rows to cluster + # rank e / 2, at dims [rank * rows_per_rank + 8 (e % 2), ...) of that rank's packed-head row. + pr_e = tx - cutlass.Int32(128) + if pr_e < num_tokens * cutlass.Int32(2): + pr_t = pr_e // cutlass.Int32(2) + pr_h = pr_e % cutlass.Int32(2) + pr_w = _ld_shared_v4_b32( + pr_stage.subview( + pr_t * cutlass.Int32(rows_per_rank) + pr_h * cutlass.Int32(8) + ).data_ptr() + ) + _st_async_v4( + _mapa_u32( + pr_row.subview( + rank * cutlass.Int32(rows_per_rank // 2) + pr_h * cutlass.Int32(4) + ).data_ptr(), + pr_t, + ), + pr_w[0], + pr_w[1], + pr_w[2], + pr_w[3], + _mapa_u32(pr_full.data_ptr(), pr_t), + ) + if warp_id == 4: + if rank < num_tokens: + # The producer reads step inputs (positions, slots) and writes the caches: ordered after the + # upstream grids by this thread's own grid wait (already satisfied: the activation landed). + prims.griddepcontrol(prims.GridDepAction.WAIT) + while not _test_wait_cluster(pr_full.data_ptr(), 0): + pass + pr_lo = pr_row.load(idx=lane * cutlass.Int32(2)) + pr_hi = pr_row.load(idx=lane * cutlass.Int32(2) + cutlass.Int32(1)) + mask = cutlass.Int32(0xFFFF) + producer_head( + rank, + n_tile, + lane, + bf16_bits_to_f32(pr_lo & mask), + bf16_bits_to_f32((pr_lo >> cutlass.Int32(16)) & mask), + bf16_bits_to_f32(pr_hi & mask), + bf16_bits_to_f32((pr_hi >> cutlass.Int32(16)) & mask), + ProducerArgs( + packed=cutlass.Int64(0), + q_out=pr_q, + index_q_out=pr_iq, + index_k=pr_ik, + kv_data=pr_kvd, + kv_scale=pr_kvs, + out_cache_loc=pr_loc, + kv_quant_scale=pr_quant, + position_ids=pr_pos, + rotary_cos_sin=pr_rope, + q_weight=pr_qw, + k_weight=pr_kw, + index_q_weight=pr_iqw, + index_k_weight=pr_ikw, + eps=pr_eps, + num_pages=pr_pages, + index_page_stride=pr_ips, + index_token_stride=pr_its, + data_page_stride=pr_dps, + data_plane_stride=pr_dls, + data_token_stride=pr_dts, + scale_page_stride=pr_sps, + scale_plane_stride=pr_sls, + ), + ) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[bool], + # prod: see m3_mx_gemv_kernel (zeros when off). + pr_q: cutlass.Int64, + pr_iq: cutlass.Int64, + pr_ik: cutlass.Int64, + pr_kvd: cutlass.Int64, + pr_kvs: cutlass.Int64, + pr_loc: cutlass.Int64, + pr_quant: cutlass.Int64, + pr_pos: cutlass.Int64, + pr_rope: cutlass.Int64, + pr_qw: cutlass.Int64, + pr_kw: cutlass.Int64, + pr_iqw: cutlass.Int64, + pr_ikw: cutlass.Int64, + pr_ips: cutlass.Int64, + pr_its: cutlass.Int64, + pr_dps: cutlass.Int64, + pr_dls: cutlass.Int64, + pr_dts: cutlass.Int64, + pr_sps: cutlass.Int64, + pr_sls: cutlass.Int64, + pr_pages: cutlass.Int32, + pr_eps: cutlass.Float32, + prod: cutlass.Constexpr[bool], +) -> None: + # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. + assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( + "x_swiglu needs x_bf16 and excludes din" + ) + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, x_swiglu, + pr_q, pr_iq, pr_ik, pr_kvd, pr_kvs, pr_loc, pr_quant, pr_pos, pr_rope, pr_qw, pr_kw, pr_iqw, pr_ikw, pr_ips, + pr_its, pr_dps, pr_dls, pr_dts, pr_sps, pr_sls, pr_pages, pr_eps, prod, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/op.py new file mode 100644 index 000000000000..efe4879f9fda --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/op.py @@ -0,0 +1,645 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) at the TP2 per-rank projection shapes (SM107). + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring[, contig, nacc, a_tmem, st_async]) of the TP2 per-rank projections. The optional fields default +# to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. SM107 graph us at M 1 / 4 / 8 over weights rotated beyond L2, against the +# stock MXFP8 quantize + FlashInfer CuTe DSL GEMM: QKV 4.11 / 4.12 / 4.10 vs 8.6-8.8, o_proj 3.90 / 3.96 / 3.70 vs +# 6.9-7.2, shared gate_up 3.11 / 3.09 / 3.10 vs 8.1-8.4, shared down 2.22 vs 4.7-4.9, dense gate_up 8.15 / 8.56 / 8.55 +# vs 10.2-10.5, dense down 5.89 / 5.72 / 5.46 vs 8.7-9.0. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + # The sm_107 tp8 target's per-rank projections (8 ranks over two trays), (split, ring) from the SM107 sweep 778596 + # (locked, MXFP8 in, back-to-back PDL, rows 1 / 4 / 8; the GB300 tp8 table has none of these shapes on SM107). + # sparse-layer QKV + index projection: 2.37 / 2.43 / 2.42 us + (1536, 6144): (8, 8), + # dense-layer QKV: 2.34 / 2.35 / 2.36 us + (1280, 6144): (8, 8), + # o_proj: 2.05 / 2.07 / 2.07 us + (6144, 1024): (2, 12), + # shared-expert gate_up: 2.16 / 2.17 / 2.17 us + (768, 6144): (8, 8), + # shared-expert down: 1.49 / 1.50 / 1.52 us + (6144, 384): (1, 8), + # dense-layer gate_up: 3.09 / 3.12 / 3.08 us + (3072, 6144): (4, 12), + # dense-layer down: 2.29 / 2.30 / 2.29 us + (6144, 1536): (2, 8), +} +# Per-token-count configs (N, K, M); none at TP2. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = {} +# Above 8 tokens (wider than the TP2 decode step); none at TP2. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = {} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +# (rule 1b, deterministic) The decode projections at up to STRIDED_MAX_ROWS rows split K over strided k-tiles per +# split rank (contig 0: rank r owns k-tiles r, r + split, ...) instead of contiguous ranges: another fixed fp32 order. +# SM107 TP8 4 rows (783579), us/call contiguous -> strided: sparse QKV 2.405 -> 2.371, dense QKV 2.323 -> 2.291, +# o_proj 2.061 -> 2.019, dense gate_up 3.047 -> 3.006, dense down 2.313 -> 2.287; deterministic (3 runs bitwise) and +# max error vs the fp64 product equal to the contiguous order's. Kill: False. +STRIDED_K_TILES = True +STRIDED_MAX_ROWS = 8 # capacity: fixed: the strided-input kernels' built rows +_STRIDED = frozenset({(1536, 6144), (1280, 6144), (6144, 1024), (3072, 6144), (6144, 1536)}) + +# #117: the producer epilogue serves the sparse-layer QKV + index projection, 12 packed heads of 128 (8 q, k, v, +# index-q, index-k), one head per split-8 cluster. +PROD_N = 12 * 128 # capacity: fixed: the producer's packed heads +# The kernel's producer arguments on a call without the producer epilogue (20 addresses / strides, pages, eps). +PROD_OFF = (0,) * 21 + (0.0,) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + if ( + STRIDED_K_TILES + and m is not None + and m <= STRIDED_MAX_ROWS + and (n, k) in _STRIDED + and len(cfg) == 2 + ): + cfg = cfg + (0,) + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # 17..32 rows at K < 1024 decline (an unsplit K below 1024 is slower than the FlashInfer GEMM there); no TP2 + # projection has K < 1024. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], + x_swiglu: bool = False, + prod: bool = False, +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + if prod and not ( + n == PROD_N + and split == 8 + and m <= split + and st_async + and not x_bf16 + and not x_swiglu + and ar is None + ): + # The producer epilogue: the sparse QKV + index projection (one 128-row packed head per cluster of 8), at most + # one token per cluster rank, the st.async exchange, MXFP8 input, no all-reduce push. + return None + if nacc != 1: + # More than one TMEM accumulator is broken on SM107 (783556 / 783557 / 783579: nondeterministic or garbage + # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + prod, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile( + key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple, prod_args: tuple +): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + prod, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + x_swiglu, + *prod_args, + prod, + ) + return fn + + +def _views( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight, + weight_scale, + y, + x_bf16: bool, + x_swiglu: bool = False, +) -> tuple: + # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. + k = x.shape[1] // 2 if x_swiglu else x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, + x_swiglu: bool = False, + prod_args: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + if x_swiglu: + k //= 2 + n = weight.shape[0] + if weight.shape[1] != k: + return None + prod = prod_args is not None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu, prod) + if resolved is None: + return None + key, runtime_ar = resolved + if not prod: + prod_args = PROD_OFF + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + # prod: no packed row is stored (the producer epilogue consumes it); y is the placeholder. + y = torch.empty(0 if prod else m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra, prod_args) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra, *prod_args) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, + x_swiglu: bool = False, + prod: bool = False, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu, prod) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros( + m, + 2 * k if x_swiglu else k, + dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, + device=dev, + ) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev), PROD_OFF) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_swiglu( + gate_up: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``m3_mx_gemv_bf16(swiglu(gate_up), ...)`` in one launch: ``gate_up`` bf16 [M <= 8, 2K] (contiguous, gate then + up) is activated in the kernel's quantize prologue with the swigluoai arithmetic of the Triton silu_and_mul kernel + (alpha 1.702, limit 7, the (up + 1) offset; bit for bit), then quantized as the bf16 path does.""" + if ( + gate_up.dtype != torch.bfloat16 + or gate_up.dim() != 2 + or not gate_up.is_contiguous() + or gate_up.data_ptr() % 16 != 0 + or gate_up.shape[1] % 2 != 0 + ): + return None + return _run(gate_up, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar, True) + + +def m3_mx_gemv_producer( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + q: torch.Tensor, + index_q: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + trigger_early: bool = True, + use_pdl: Optional[bool] = None, +) -> bool: + """#117: ``m3_producer(m3_mx_gemv(xq, xsf, weight, weight_scale), ...)`` in one launch, bit for bit: the QKV + + index GEMV of the sparse layer (``weight`` [12 x 128, K]) whose owner epilogue hands each token's packed head to + the producer phase (``m3_layer.phase_producer.producer_head``). Writes ``q`` FP8 [T, 8, 128] and ``index_q`` FP8 + [T, 1, 128] and the token's cache slot (NVFP4 K / V + scales, FP8 index-K) exactly as the producer does; the + caches and tensors are those ``minimax_m3_producer_pw`` takes (kv_data [pages, 2, 1, 128, 64] u8, kv_scale [pages, + 2, 1, 128, 8] u8, index_k [pages, 1, 128, 128] e4m3). Returns False (nothing launched) when the call is declined. + """ + m = xq.shape[0] + if ( + m > 8 + or weight.shape[0] != PROD_N + or q.shape != (m, 8, 128) + or index_q.shape != (m, 1, 128) + or not (q.is_contiguous() and index_q.is_contiguous()) + or kv_data.dim() != 5 + or kv_scale.dim() != 5 + or index_k.dim() != 4 + ): + return False + tensors = ( + q, + index_q, + index_k, + kv_data, + kv_scale, + out_cache_loc, + kv_quant_scale, + position_ids, + rotary_cos_sin, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + ) + prod_args = ( + *(t.data_ptr() for t in tensors), + index_k.stride(0), + index_k.stride(2), + kv_data.stride(0), + kv_data.stride(1), + kv_data.stride(3), + kv_scale.stride(0), + kv_scale.stride(1), + kv_data.shape[0], + float(eps), + ) + out = _run( + xq, xsf, weight, weight_scale, False, trigger_early, None, use_pdl, None, False, prod_args + ) + return out is not None + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode.py new file mode 100644 index 000000000000..e3694274c124 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode.py @@ -0,0 +1,1754 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +GB300/SM103 tuning targets GQA group 16, top-k 16 and four-token Eagle3 +verification. Split-K, query pairing and launch tables require four KV heads; +the linear-softmax policy also covers one and two KV heads. Other hardware +and shapes retain the default policies, including supported SM100 NVFP4. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* A packed NVFP4 cache is read in place, dequantized in-register from its E2M1 + nibbles and E4M3 block scales. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* A single top-k chunk writes the final output directly, avoiding scratch and + the otherwise redundant LSE merge. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. + +This is the MiniMax-M3 NVFP4 / sm_107 / tp8 target's copy. On top of the +in-tree decode it routes long-context TP8 steps (8 query heads, one KV head) to the one-launch +``m3_attn_lc8`` kernel, short-sequence verify steps of at least +``MQ_MIN_TOKENS`` rows to the multi-query producer, releases PDL dependents +early, and has the merge emit the MXFP8 o_proj hand-off (``m3_prequant.SLOT``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + check_decode_span_shape, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +from . import m3_attn_lc8, m3_prequant, m3_report + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +_TL_DTYPES = { + torch.bfloat16: tl.bfloat16, + torch.float16: tl.float16, + torch.float32: tl.float32, +} + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size + +# NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale +# byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. +NVFP4_SF_VEC_SIZE = 16 + + +# Fewest query rows for which the merge writes the MXFP8 o_proj hand-off: every decode step, so the MXFP8 o_proj +# never runs its own quantize kernel after attention. +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# The sparse decode (LC8, the Triton merge) also writes the MXFP8 o_proj hand-off (bit for bit +# torch.ops.trtllm.mxfp8_quantize of its BF16 rows) and publishes it to m3_prequant.SLOT. +MERGE_MXFP8_HANDOFF = True +# The one-launch long-context decode (m3_attn_lc8) at the TP8 heads: its shapes (modeling's SPARSE_DECODE_LC8 keys: +# "prewait" = up to m3_attn_lc8.PREWAIT_MAX_ROWS query rows, "plain" above) and its report path. Short-sequence +# identity steps run it too, taking the identity selection from the device lengths (LC8_IDENTITY, with its merge-once +# form). +LC8_KEYS = frozenset({"plain", "prewait"}) +LC8_PATH = "m3_sparse_decode_lc8" +LC8_IDENTITY = True +# The identity-selection producer's shapes (modeling's SPARSE_DECODE_PRODUCER keys) and its report path. +PRODUCER_KEYS = frozenset({"multi_query", "per_row"}) +PRODUCER_PATH = "sparse_decode_producer" + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and swizzled UE8M0 + scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips its own + quantization. One 128-row buffer per width serves every layer: each o_proj consumes it before the next layer's + merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + +MERGE_MXFP8 = MergeMxfp8() + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + if IDENTITY: + # Slot `it` of the identity selection, bounded by this row's own kv extent above. + blk = it + else: + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, + # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + if IDENTITY: + blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) + blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) + else: + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for cand in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + if IDENTITY: + blk_0 = tl.where(valid_0, idx_0, sentinel) + blk_1 = tl.where(valid_1, idx_1, sentinel) + else: + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the + # pool, reads nothing. + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. Because this trigger precedes this grid's own + # wait, it says nothing about what upstream has finished: a dependent may read only step-invariant data + # (weights) before its griddepcontrol.wait, never this step's attention output, q, top-k table or lengths. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) +# Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) + + +# Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query +# tokens; below it the per-token producer is as fast. +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover +# Warps per multi-query producer program. +MQ_NUM_WARPS = 8 + + +def _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + identity=False, + identity_step=False, +) -> None: + """minimax_m3_sparse_attn_decode through m3_attn_lc8: same table, output and MXFP8 hand-off contract. ``identity`` + takes the short-sequence identity selection from the device lengths (the table is not read); ``identity_step`` + marks a step with no select (the kernel reads q after its grid wait), also when the staged table is read.""" + use_pdl = get_env_enable_pdl() + total_q, num_heads, head_dim = q.shape + mxfp8 = None + m3_prequant.SLOT.clear() + if MERGE_MXFP8_HANDOFF and MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(num_heads * head_dim, output.device) + # The early trigger lets the o_proj launch and stream its weights while attention runs. + m3_attn_lc8.m3_lc8_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + linear_softmax=_sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=1, + gqa_group_size=num_heads, + max_topk=int(topk_idx.shape[-1]), + decode_query_len=decode_query_len, + ), + mxfp8=mxfp8, + use_pdl=use_pdl, + identity=identity, + identity_step=identity_step, + prefetch=m3_attn_lc8.PREFETCH.take(), + next_kv=m3_attn_lc8.NEXT_KV.take(), + ) + if m3_attn_lc8.ONE_DIV_MERGE and m3_attn_lc8.MERGE_ONCE: + m3_report.taken(ONE_DIV_MERGE_PATH, total_q) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), mxfp8[0][:total_q], mxfp8[1] + ) + + +# Engagement path of the identity-selection producers reading their blocks from the kv lengths (IDENTITY). +DEVICE_IDENTITY_PATH = "sparse_decode_device_identity" +# Engagement path of the LC's one-division split-K merge (MSA_ONE_DIV_MERGE). +ONE_DIV_MERGE_PATH = "sparse_decode_lc_one_div_merge" + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + q may be FP8, as a fused QK-norm/RoPE producer emits it. The kernel widens + it to output.dtype in-register rather than making the caller materialize a + widened copy, which would cost a standalone kernel on every sparse layer. + The attention math is unchanged either way, since E4M3 widens exactly. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. + + Passing the four NVFP4 scale tensors selects the packed-NVFP4 cache instead: + k_paged/v_paged then hold two E2M1 elements per byte, the block scales are + paged E4M3 bytes in the layouts msa_scatter writes, and the two per-tensor + dequant scales are single-element FP32 device tensors. + + num_topk_chunks overrides the split-K factor and exists for tests: the + merged result must not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) + nvfp4_args = (k_block_scale, v_block_scale, k_global_scale, v_global_scale) + kv_nvfp4 = any(arg is not None for arg in nvfp4_args) + scale_cols = head_dim // NVFP4_SF_VEC_SIZE + if kv_nvfp4: + capability = torch.cuda.get_device_capability(q.device) + if capability not in _SM100F_CAPABILITIES: + raise NotImplementedError( + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; " + f"got SM{capability[0] * 10 + capability[1]}." + ) + _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) + # Two E2M1 elements per byte; Triton must see them unsigned so that the + # high-nibble shift does not sign-extend. + k_paged, v_paged = k_paged.view(torch.uint8), v_paged.view(torch.uint8) + k_block_scale, v_block_scale = ( + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused arguments. + scale_arg = kv_scale if use_scale else output + k_sc_arg, v_sc_arg = (k_block_scale, v_block_scale) if kv_nvfp4 else (output, output) + k_gs_arg, v_gs_arg = (k_global_scale, v_global_scale) if kv_nvfp4 else (output, output) + + widen_q = q.dtype in _FP8_DTYPES + if widen_q and output.dtype not in _TL_DTYPES: + raise ValueError( + f"MiniMax-M3 sparse decode cannot widen FP8 q into {output.dtype}; " + f"supported compute dtypes are {sorted(d.__str__() for d in _TL_DTYPES)}." + ) + # Packed NVFP4 dequantizes to FP16. The production producer supplies FP8 + # q, which widens exactly to FP16, so keep both dot operands there instead + # of converting every K/V tile to BF16. Accumulation remains FP32. + if widen_q: + compute_dtype = tl.float16 if kv_nvfp4 else _TL_DTYPES[output.dtype] + else: + compute_dtype = tl.float32 + + # Long-context TP8 steps: attention, split-K merge and MXFP8 hand-off in one launch. A selected table on an NVFP4 + # cache with FP8 q at the TP8 heads (8 query heads, one KV head) is the kernel's (LC8_KEYS: "prewait" up to + # m3_attn_lc8.PREWAIT_MAX_ROWS query rows); operands it declines are a miss. + lc_key = "prewait" if total_q <= m3_attn_lc8.PREWAIT_MAX_ROWS else "plain" + lc_covered = ( + kv_nvfp4 + and widen_q + and (not identity_blocks or (LC8_IDENTITY and m3_attn_lc8.MERGE_ONCE)) + and num_topk_chunks is None + and q.dim() == 3 + and int(q.shape[1]) == m3_attn_lc8.NUM_HEADS + and num_kv_heads == 1 + and lc_key in LC8_KEYS + ) + if lc_covered and not m3_attn_lc8.lc8_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ): + m3_report.missed( + LC8_PATH, + total_q, + "the step's q / KV pages / scales / top-k table / output are not the layout the kernel is built for", + ) + lc_covered = False + if lc_covered: + _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + # The identity selection from the device lengths (the staged table is not read). + identity=identity_blocks, + # Every identity step (no select ran): q after the grid wait. + identity_step=identity_blocks, + ) + m3_report.taken(LC8_PATH, total_q) + return + + query_group_size = ( + _sm100f_nvfp4_query_group_size( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else 1 + ) + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = ( + _sm100f_nvfp4_num_topk_chunks( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else None + ) + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one + # program per (request, block) covers all of them and dequantizes each block once. + use_mq = ( + identity_blocks + and total_q >= MQ_MIN_TOKENS + and kv_nvfp4 + and num_kv_heads == 1 + and decode_query_len > 1 + and head_dim == 128 + and gqa_group_size * decode_query_len <= 128 + ) + if identity_blocks and kv_nvfp4 and widen_q: + # The identity-selection route is modeling's SPARSE_DECODE_PRODUCER ("per_row"; "multi_query" needs one KV + # head, so it never runs at the TP2 heads); a variant the range does not certify is a miss (raises under + # require). + producer_key = "multi_query" if use_mq else "per_row" + if producer_key not in PRODUCER_KEYS: + m3_report.missed( + PRODUCER_PATH, + total_q, + f"the identity-selection producer variant {producer_key!r} is not certified at {num_kv_heads} KV heads", + ) + if use_mq: + num_topk_chunks = max_topk + direct_output = num_topk_chunks == 1 + if direct_output: + # With one chunk, acc_o is already normalized and final. Point the + # producer at the caller's output and compile away the LSE store. + # Keeping a real dummy pointer preserves Triton's pointer contract. + o_arg = output + lse_arg = output + o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) + lse_strides = (0, 0, 0) + else: + # Persistent arena rather than torch.empty, so the partials keep one + # address across CUDA graph replays. FP32 keeps the split-K factor from + # perturbing the merged result. + reserve = torch.cuda.is_current_stream_capturing() + o_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + o_strides = tuple(o_arg.stride()) + lse_strides = tuple(lse_arg.stride()) + + use_pdl = get_env_enable_pdl() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = ( + _sm100f_nvfp4_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + if total_q % query_group_size or decode_query_len % query_group_size: + raise ValueError( + f"query group {query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." + ) + + producer_kernel = ( + _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel + ) + device_identity = identity_blocks + if device_identity and not use_mq: + m3_report.taken(DEVICE_IDENTITY_PATH, total_q) + single_kernel_args = ( + { + # No split-K arrival counters: the separate merge kernel combines the chunks. + "merge": (output, None), + "tail": (output.stride(0), output.stride(1), output.stride(2), num_kv_heads), + # An early release pays off when the dependent is the consumer of the + # final output (it can prefetch), or a merge grid that itself releases + # its dependent at entry: the o_proj then launches during this kernel and + # streams its weights while attention runs. Above 8 query rows the + # concurrent weight stream slows the larger attention grid more than the + # prefetch saves (GB300 TP4: -1.1% at 48 rows). + "options": {"FUSED_MERGE": False, "EARLY_TRIGGER": direct_output or total_q <= 8}, + } + if query_group_size == 1 + else {"merge": (), "tail": (), "options": {}} + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + else {} + ) + if use_mq: + from .triton_sparse_decode_mq import launch_mq_producer + + launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + sm_scale=sm_scale, + decode_query_len=decode_query_len, + use_pdl=use_pdl, + early_trigger=False, + num_warps=MQ_NUM_WARPS, + ) + else: + producer_kernel[ + (total_q // query_group_size) * num_topk_chunks, + num_kv_heads, + ]( + q, + k_paged, + v_paged, + scale_arg, + k_sc_arg, + v_sc_arg, + k_gs_arg, + v_gs_arg, + topk_idx, + o_arg, + lse_arg, + *single_kernel_args["merge"], + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + k_sc_arg.stride(0), + k_sc_arg.stride(1), + v_sc_arg.stride(0), + v_sc_arg.stride(1), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + *o_strides, + *lse_strides, + block_table.stride(0), + block_table.shape[1], + k_paged.shape[0], + *single_kernel_args["tail"], + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + WIDEN_Q=widen_q, + COMPUTE_DTYPE=compute_dtype, + USE_PDL=use_pdl, + KV_NVFP4=kv_nvfp4, + USE_LINEAR_SOFTMAX=use_linear_softmax, + SCALE_COLS=scale_cols, + DIRECT_OUTPUT=direct_output, + # Identity selection from the lengths this kernel reads, not the staged table. + IDENTITY=device_identity, + **single_kernel_args["options"], + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + m3_prequant.SLOT.clear() + if not direct_output: + width = num_heads * head_dim + mxfp8 = ( + MERGE_MXFP8.buffers(width, output.device) + if MERGE_MXFP8_HANDOFF + and MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS + and head_dim == 128 + and output.dim() == 3 + and output.is_contiguous() + and output.dtype == torch.bfloat16 + else None + ) + mxfp8_args = ( + { + "q8_ptr": mxfp8[0], + "sf_ptr": mxfp8[1], + "num_sf_k_tiles": width // 128, + "MXFP8_OUT": True, + } + if mxfp8 is not None + else {} + ) + merge_launch_options = ( + _sm100f_nvfp4_merge_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_arg, + lse_arg, + output, + head_dim, + *o_strides, + *lse_strides, + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + # Release the o_proj at the merge's entry so it streams its weights during the merge. + EARLY_TRIGGER=True, + **mxfp8_args, + **pdl_launch, + **merge_launch_options, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output.view(total_q, width), mxfp8[0][:total_q], mxfp8[1]) + + +__all__ = [ + "MERGE_MXFP8", + "MERGE_MXFP8_MIN_TOKENS", + "MergeMxfp8", + "NVFP4_SF_VEC_SIZE", + "SPARSE_BLOCK_SIZE", + "_sm100f_nvfp4_merge_launch_options", + "_sm100f_nvfp4_num_topk_chunks", + "_sm100f_nvfp4_query_group_size", + "_sm100f_nvfp4_use_linear_softmax", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__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_107__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py new file mode 100644 index 000000000000..9a279b669def --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py @@ -0,0 +1,6151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_107 / tp8 -- self-contained modeling code. + +MiniMax-M3's language model: 60 layers, hidden 6144, 64 query heads and 4 KV heads of dim 128. Layers 0-2 are dense +attention with a dense SwiGLU-OAI MLP; layers 3-59 are MSA block-sparse attention (an index branch scores 128-token +blocks and the top 16 are attended) with a MoE of 128 routed experts (top 4, sigmoid scores with a correction bias, +renormalized and scaled by 2) plus one shared expert. The checkpoint is a modelopt mixed-precision export: NVFP4 routed +experts, MXFP8 projections, fp32 router, bf16 embedding and head; the KV cache is NVFP4 for the sparse layers and FP8 +for the dense ones. + +``tp8`` is eight Rubin GPUs over two trays (multi-node NVLink): head-split attention (8 query heads, one KV head and +one index head per rank, each KV head shared by two ranks) and tensor-parallel experts (every expert on every rank, +intermediate 384 per rank), no attention DP. +Speculative decoding, when configured, is one-model Eagle3 (draft length 3): the BF16 drafter (``M3Eagle3Draft``) and +the BF16 target head, both at checkpoint precision; decode steps of up to 2 requests verify and draft through fused +kernels at that precision (``M3Eagle3Worker``). + +**The core** (``MiniMaxM3Tp8ModelingV2``) holds the weights (``weights.MODEL_WEIGHTS``) and the runtime resources the +steps use -- one MSA attention backend per layer, the TP all-reduce communicator, the RoPE table, the MXFP8 GEMM +dispatch state the engine tunes -- and dispatches each step, on the host, from its context count alone +(``_select_target``): + +* ``DecodeTarget`` takes every pure generation step. Up to ``DECODE_MAX_TOKENS`` tokens (up to 2 requests, each one + token or a draft verify of 1 + 3 rows) with 1-D input ids it runs the decode-step kernels and picks each layer + boundary's collective by the step's token bucket (``DECODE_PLANS``); it runs its own generic layer for the first + layer when the fused embedding does not take the step, and for every layer of a wider step. +* ``PrefillTarget`` takes every step with context rows: prefill chunks and mixed steps, every layer in its generic + form over the MSA backend. + +The two targets share no body: each owns its layer, attention and MoE compositions and its ``step_args``, the one +projection of the step metadata its forward reads. What both call are the module-level op adapters, each picking the +kernel for one call from shapes and layouts, which read no step metadata themselves. + +Every fast path the decode step can take, the rows each covers (``FAST_PATHS``, each paired with the entry whose +``CELLS`` certify it) and what a step that cannot take one does are below: **uncovered** (outside every covered +range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise +``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). A row range is empty unless a +landed cell has certified it at the TP8 geometry on Rubin. Every fast path a GEN rank takes prints a rank-tagged +engagement line once per (path, rows) (``taken``). Kernel modules report through ``kernels/m3_report``, which this +module wires to the same policy at import. + +**Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the +next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the +long-context decode for the o_proj. It publishes it in ``kernels/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +MXFP8 projection on a decode graph takes it, and every take clears the slot. The decode boundaries' push targets (the +o_proj / dense-down / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's exchange) +are passed explicitly from the layer to the producer and back. + +Compiled kernels (the MoE engine, the decode GEMVs with their push variants, the index scorer, the long-row select, +the full-context decode plans) are built before the executor's warmup; a forward never compiles one, and one it finds +missing is a miss. Other state a fast path needs (multicast workspaces, push descriptors) is built on the eager decode +steps; a CUDA-graph capture only reads it, so a capture that finds it missing is a miss (``capture_missed``). + +**What this target asserts rather than adapts**: SM 10.7, the topology and checkpoint (``routing.py``); the +MiniMax-M3 MSA sparse attention runtime with the fused five-way QKV/index projection, an FP8 indexer and an FP8 or +NVFP4 KV cache, no torch.compile and no piecewise prefill graphs (``_validate_runtime``); text-only input (the vision +tower is not built and a step carrying image input raises). +""" + +from __future__ import annotations + +import atexit +import copy +import dataclasses +import os +import threading +import time +from dataclasses import dataclass +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + resolved_mode, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, +) +from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + _ensure_draft_vocab_size, + _get_requested_draft_moe_backend, + _set_draft_kv_cache_quant_algo, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + TensorParallelMode, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen +from tensorrt_llm._utils import mpi_rank, prefer_pinned +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import weights as _weights +from .kernels import m3_draft_step as m3_draft_step_kernel +from .kernels import ( + m3_drafter_attn, + m3_fc_wait, + m3_fold_tail, + m3_mc_allreduce, + m3_prequant, + m3_report, + m3_spec_tail, +) +from .kernels import m3_fullctx_decode as m3_fullctx_decode_kernel +from .kernels import m3_index_score as _index_score_kernel +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 import triton_sparse_decode as _sparse_decode_kernel +from .kernels.long_select import long_select_blocks +from .kernels.mx_gemv import op as mx_gemv +from .ops.activation.m3_silu_and_mul import m3_silu_and_mul +from .ops.attention.m3_draft_attention import m3_draft_attention +from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope +from .ops.attention.m3_fullctx_decode import m3_fullctx_decode +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_msa_attention import m3_msa_attention +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 +from .ops.attention.m3_qkv_insert import m3_qkv_insert +from .ops.attention.m3_select_blocks import m3_select_blocks +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_select_long_dsl import m3_select_long_dsl +from .ops.attention.m3_sparse_decode import m3_sparse_decode +from .ops.attention.m3_sparse_decode_lc8 import m3_sparse_decode_lc8 +from .ops.comm.m3_allreduce import m3_allreduce +from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.embedding.m3_embedding import m3_embedding +from .ops.gemm.m3_bf16_linear import m3_bf16_linear +from .ops.gemm.m3_dense_swiglu import ALPHA as DENSE_ALPHA +from .ops.gemm.m3_dense_swiglu import BETA as DENSE_BETA +from .ops.gemm.m3_dense_swiglu import LIMIT as DENSE_LIMIT +from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu +from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS +from .ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for +from .ops.gemm.m3_fp32_linear import m3_fp32_linear +from .ops.gemm.m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + m3_mx_gemv, +) +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .ops.spec.m3_capture_norm import m3_capture_norm +from .ops.spec.m3_chunk_norm import m3_chunk_norm +from .ops.spec.m3_draft_input import m3_draft_input +from .ops.spec.m3_draft_silu import m3_draft_silu +from .ops.spec.m3_draft_step import m3_draft_step +from .ops.spec.m3_force_vec import m3_force_vec, m3_force_vec_req +from .ops.spec.m3_sampled_accept import m3_sampled_accept +from .ops.spec.m3_shard_accept import m3_shard_accept +from .ops.spec.m3_spec_offsets import m3_spec_offsets +from .ops.spec.m3_step_prologue import m3_step_prologue + +_SM = (10, 7) +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The largest generation step the decode target takes: 2 requests verifying 1 + 3 drafts (the GEN worker's batch). +DECODE_MAX_TOKENS = 8 # capacity: fixed: the decode kernels' step bound + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 tp4 measurements in the comments).""" + + name: str + max_tokens: int + # The attention boundary's multicast all-reduce + residual + RMSNorm (5.05 vs 5.27 us native at 1 token). + attn_mc_ar: bool + # The same kernel at the MoE boundary with the shared + routed add folded in: wins up to 48 tokens. + moe_mc_ar: bool + # The MoE-boundary all-reduce also emits the next QKV activation in MXFP8 for the decode GEMV / GEMM. + qkv_prequant: bool + + +# tiny: every step up to 8 tokens (c1..c8 without speculation, the c1 verify, every draft pass); small: the c2 +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (kernels/m3_mc_allreduce.BUCKETS). +# The multicast all-reduce + residual + RMSNorm at both boundaries of every TP2 decode step (one 8-token bucket), the +# MoE boundary's also publishing the next QKV activation in MXFP8. SM107 TP2 hop: 3.5..4.4 us at 1..8 rows against +# 8.7..9.5 us for the stock AUTO (NCCL) all-reduce + fused residual RMSNorm, and 13.3..14.0 us at the MoE boundary +# (stock add + all-reduce + mxfp8_quantize); bitwise equal to stock (ops/comm/m3_mc_allreduce_norm). +# False puts both boundaries back on the built-in all-reduce. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +MC_AR_TP2 = True +DECODE_PLANS = ( + DecodePlan( + "tiny", DECODE_MAX_TOKENS, attn_mc_ar=MC_AR_TP2, moe_mc_ar=MC_AR_TP2, qkv_prequant=MC_AR_TP2 + ), +) + +# The dense layers' decode attention (trtllm-gen, msa_decode.dense_decode_max_seq_len): the max_seq_len it is handed +# per request count. It only sizes the multi-CTA KV split (every CTA still covers its share of the actual length; the +# result changes only by the split-merge order), and the model maximum over-splits AgentX-length caches at one +# request: GB300 TP4, 16 heads, FP8 KV, 64K context: 15.1 us at 16K vs 28.6 us at 1M. Two requests are neutral (the +# grid is unchanged) and keep the step's bound. + + +# TP2 skeleton: no cap (the built-in split) until a TP2 cap is measured and certified. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) + +# The dense layers' decode attention on PrimTS split-KV decode (kernels/m3_fullctx_decode) instead of the trtllm-gen +# call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the +# BF16 merge. At TP8 (8 query heads, 1 KV head per rank) the stock split is capped at 16 cluster CTAs per (request, KV +# head), i.e. 16 of 212 CTAs at one request. Rubin TP8 geometry, FP8 KV, chained us per call at 165k / 323k, every +# recipe passing the fp32 oracle (microbenches tp8s 777734 and tp4k 777845, locked clock, cold L2): (4, 1) 12.65 / +# 15.88 at tile 16 / 96 splits / 12 K/V stages; (4, 2) 15.85 / 21.32 at tile 32 / 96 splits. At one request, 108 or +# more splits at tile 16 fall off a cliff at 165k, and fewer K/V stages are slower at both lengths. The drafter's +# one-row passes (the 8-row Q tile has no separate-reduction profile with FP16 partials, so tile 16): (1, 1) 10.75 us +# at 108 splits / 12 K/V stages (tp8b 777783); (1, 2) 13.35 at 64 splits / 12 K/V stages (tp8c 779066). +FULLCTX_DECODE_RECIPES: Mapping[Tuple[int, int], Mapping[str, Any]] = MappingProxyType( + { + (4, 1): MappingProxyType( + { + "config": {"tile_size_q": 16, "groups_tokens_heads_q": True, "kv_stages": 12}, + "splits": 96, + } + ), + (4, 2): MappingProxyType( + {"config": {"tile_size_q": 32, "groups_tokens_heads_q": True}, "splits": 96} + ), + (1, 1): MappingProxyType({"config": {"tile_size_q": 16, "kv_stages": 12}, "splits": 108}), + (1, 2): MappingProxyType({"config": {"tile_size_q": 16, "kv_stages": 12}, "splits": 64}), + } +) +# False puts the dense layers back on the trtllm-gen call at every shape. +FULLCTX_DECODE_ENABLED = True +# The dense layers' FP8 q/k/v (decode steps of the hybrid cache): the TP2 op rounds Q/K to BF16 and then to E4M3 in +# its store, the bits of the BF16 fused_qk_norm_rope output cast to FP8. False runs that BF16 kernel and a separate cast +# (the same bits, one more kernel). +DENSE_QKV_FP8_VIA_BF16 = True +# On a full-context decode step that op also stores the new K/V rows into the layer's paged pool (the bytes +# write_msa_phase_kv's indexed copies store, one launch instead of about ten). False writes them with +# write_msa_phase_kv. +DENSE_KV_INSERT_FUSED = True +# The engagement line of that fused store (dispatch.taken, one per row count and rank). +DENSE_KV_INSERT_PATH = "m3_dense_kv_insert" +# The drafter MLP's C++ SwiGLU (DRAFT_SILU_CPP, spec_worker._m3_draft_mlp): its own taken / not-taken lines. +DRAFT_SILU_PATH = "m3_draft_silu" +# The drafter boundary all-reduce fed by its projection's push (DRAFT_AR_PUSH): its own taken / not-taken lines. +DRAFT_AR_PUSH_PATH = "m3_draft_ar_push" +# The split-KV merge on the BF16 op (minimax_m3_sm107_tp2_fullctx_reduce) instead of PrimTS's FP16 reducer and a cast: +# Rubin, 165k, graph-us per call (1, 1) 14.5 vs 18.5, (4, 1) 17.3 vs 21.3, (1, 2) 18.5 vs 23.1, (4, 2) 23.3 vs 27.5. +FULLCTX_DECODE_OWN_REDUCER = True +# The TP8 merge also writes the MXFP8 copy of its BF16 rows (E4M3 + 128x4-swizzled UE8M0, bit for bit +# trtllm::mxfp8_quantize) into the MERGE_MXFP8 buffers and publishes it to m3_prequant.SLOT, so the dense layer's +# MXFP8 o_proj skips its own quantization. Kill: False leaves the o_proj to quantize its BF16 input. +FULLCTX_DECODE_MXFP8 = True +FULLCTX_DECODE_MXFP8_PATH = "m3_fullctx_merge_mxfp8" +# The K/V length bound the full-context plans are compiled for; a step whose cache length bound is longer is a miss. It +# covers the TP2 AgentX config's max_seq_len (1 << 20) plus the speculative rows the cache reserves past it. +FULLCTX_DECODE_MAX_KV_LEN = 1 << 21 # capacity: fixed: the full-context plans' compiled K/V bound +# The Eagle3 drafter layer's generation phase on the same full-context decode (kernels/m3_drafter_attn: the in-tree +# RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its +# split at 8 CTAs per (request, KV head) on SM 10.7: 45-46 us per draft pass at 165k. +DRAFTER_FULLCTX_ENABLED = True +# On those passes the drafter's RoPE, FP8 query and FP8 K/V append run as one launch (the op +# minimax_m3_sm107tp8_drafter_rope_fp8_kv: the in-tree preprocess's GPT-NeoX rotation and unit-scale E4M3 store, and +# the torch E4M3 cast of its BF16 query) instead of thop.trtllm_gen_generation_preprocess and the cast. Kill constant. +DRAFTER_ROPE_FUSED = True +# The engagement line of that launch (dispatch.taken, one per row count and rank). +DRAFTER_ROPE_PATH = "m3_drafter_rope" +# thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request +# multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by +# that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, +# token-0 drafts). kernels/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# calls: what the build writes (padded rows included), so bitwise the op with a fresh map. False:. +PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" +# How the rotation's products are summed, as the in-tree preprocess build sums them (0: two products then the add, +# 1: fma(cos, x, sin * pair), 2: fma(sin, pair, cos * x)); the op's bitwise test pins it. +DRAFTER_ROPE_MODE = 1 + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the index scorer launches and streams its pre-wait index-K prefetch while the producer +# runs (its outputs are read only after the scorer's own grid wait). Rubin TP2, 57-layer PDL chain behind a 21 MB +# pre-wait L2 prefetch (the 165K scorer's shape): 3.59 (native) -> 1.93 us per layer at 4 rows, 3.64 -> 2.00 at 8. +PRODUCER_EARLY_TRIGGER = True +# The producer also reads its per-step inputs (position ids, cache slots, RoPE rows) before its grid wait, leaving only +# the packed QKV row after it. Their writers are host-to-device copies before the replay and plain-launched kernels +# earlier in the step (the input preprocessing, msa_kv_update or its torch fallback), so they are final before the +# producer can start; a PDL-launched writer of any of them would have to turn this off. The call site's +# position_ids.reshape(-1).contiguous().to(torch.int32) and packed.contiguous() launch a plain torch copy right before +# the producer when the dtype or layout differs: also a plain writer, so the contract holds. Bit-identical outputs. +# TP8: True runs this target's copy of the producer (trtllm::minimax_m3_sm107_tp8_producer_pw, +# minimaxM3Sm107Tp8/producerPw8.cu, the TP2 producerPw2 at 8 query heads) with the pre-wait reads; False (the kill +# value) runs the shared minimax_m3_producer_pw at 8 query heads, which reads them after its wait. +# Step-input writers added since: the step-input prologue's C++ op (STEP_PACK_CPP_PROLOGUE) and the C++ spec fold tail +# (STEP_FOLD_TAIL_CPP) are plain launches like the Triton kernels they replace (their tests assert it on the captured +# graph); STEP_KV_LENS_UPDATE_ONCE only drops an earlier writer whose outputs the surviving msa_kv_update rewrites. +# The spec step correction (STEP_SPEC_OFFSETS_ONE_LAUNCH: position_ids, kv_lens) is a plain launch too (its test +# asserts it). +PRODUCER_PREWAIT_STEP_INPUTS = True +# #117: on a decode step of at most 8 rows the sparse layer's QKV + index GEMV runs the producer in its owner epilogue +# (kernels/mx_gemv m3_mx_gemv_producer, m3_layer.phase_producer.producer_head): one launch and one PDL edge fewer per +# sparse layer, bit for bit the GEMV + m3_producer chain (q, index-q and every cache byte). The QKV linear defers its +# GEMV (``_linear``'s ``defer``: a ``_DeferredGemv``) and the producer call site runs the fused kernel, or the plain +# GEMV when it declines. +# Edges bench 785463 (TP8 c1, 57 layers, cold L2, both orders): -0.61 .. -0.98 us/layer at 165K / 323K. Kill: False. +QKV_PRODUCER_FOLD = True +QKV_PRODUCER_FOLD_PATH = "m3_qkv_producer_fold" + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. SM107: 1.55..1.65 us against +# 3.4..6.0 us for the stock embedding + RMSNorm at 1..8 rows (the gather bitwise, the norm one ulp off on ~5e-6 of the +# elements: summation order); it also hands the first layer its residual, +# so layer 0 takes the decode layer instead of the generic one. False runs layer 0 on the generic layer. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +EMBED_FIRSTNORM_TP2 = True +EMBED_FIRSTNORM_MAX_TOKENS = DECODE_MAX_TOKENS if EMBED_FIRSTNORM_TP2 else 0 +# The o_proj pushes its rows into the attention boundary's multicast all-reduce (prepushed mode 2) up to this many +# tokens; wider steps run the GEMV and a plain multicast all-reduce. SM107 TP2: -0.5..-1.2 us per layer, bitwise equal +# to the unpushed boundary. False keeps the GEMV and the boundary separate. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +ATTN_PUSH_TP2 = True +ATTN_PUSH_MAX_TOKENS = DECODE_MAX_TOKENS if (MC_AR_TP2 and ATTN_PUSH_TP2) else 0 +# The dense layers' MLP down projection pushes its rows into the post-MLP boundary's multicast all-reduce (prepushed +# mode 2), as the o_proj does into the attention boundary: the all-reduce no longer waits for the whole down GEMV grid +# before loading and pushing the rows itself. SM107 TP2 (stack-ON nsys 773230): the two non-capture dense layers' +# post-MLP hops end 5.5-5.8 us after the down GEMV vs 2.4-2.5 us for a pushed hop. Bitwise equal to the unpushed +# boundary (the same bf16 words). False keeps the down GEMV and the boundary separate. +DENSE_PUSH_TP2 = True +DENSE_PUSH_MAX_TOKENS = DECODE_MAX_TOKENS if (MC_AR_TP2 and DENSE_PUSH_TP2) else 0 +# The generic layer's boundaries (prefill, mixed steps, Eagle3 capture layers): the most rows that take the multicast +# all-reduce + residual + RMSNorm, alone and with the shared + routed add folded in: the decode boundaries' certified +# rows, under the same constant (MC_AR_TP2; off = the built-in all-reduce). +GENERIC_MC_AR_MAX_TOKENS = DECODE_MAX_TOKENS if MC_AR_TP2 else 0 +GENERIC_MC_AR_FOLDED_MAX_TOKENS = DECODE_MAX_TOKENS if MC_AR_TP2 else 0 +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch, ops/moe/m3_moe_engine at +# the TP2 width 1536). Rubin TP2 ablation chain (57 layers behind a 58.4 MB PDL producer, row = (P + MoE) - P, both +# orders): 31.3 / 50.7 us/layer at 4 / 8 tokens against the built-in MoE's 42.4 / 55.4 (FlashInfer MXFP8 shared +# expert); bit for bit the built-in block's shared.add_(routed). KILL_MOE_ENGINE: the built-in MoE at every row count. +KILL_MOE_ENGINE = False +MOE_ENGINE_MAX_TOKENS = 0 if KILL_MOE_ENGINE else 8 +# The engine's push of its sum into the MoE boundary's multicast all-reduce (prepushed mode 2) on the rows the +# engine serves, with the multicast boundary (DECODE_PLANS' moe_mc_ar): the boundary reads the pushed words instead +# of loading and pushing the rows itself. The forward arms the push only on rows MOE_PUSH covers. +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS if MC_AR_TP2 else 0 +# When the engine pushes its sum into the boundary all-reduce, every CTA releases the PDL dependents right after its +# grid dependency wait instead of after the push, so the all-reduce (which polls the pushed tagged words, not the grid) +# becomes resident on the SMs the engine frees. Scheduling only: the same bits (world-8 push matrix). Off at TP8: Rubin +# TP8 greedy early column, same-job A,B + B,A, flat (period +20..+26 at 155K..323K); True turns it on. +MOE_EARLY_AR_RELEASE = False +# A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel +# over the target lm_head's BF16 shard logits (kernels/m3_shard_accept), at the configured draft length. +SHARD_ACCEPT_DRAFT_LEN = 3 # capacity: fixed: the built Eagle3 draft length +# 0: (the gathered logits + argmax run). Rubin TP2 greedy verify (microbench): stock 120.4 us at b1; fused +# with the coded head 117.0 us; the BF16-head + shard-accept arm is not measured yet. +SHARD_ACCEPT_MAX_REQUESTS = 2 # capacity: fixed: the TP2 GEN batch +# A sampled (temperature + top-p, advanced_sampling_mode no_topk) Eagle3 verify step of up to this many requests +# verifies with the fused sampled acceptance (kernels/m3_sampled_accept, the TP2-shard kernel) over the target +# lm_head's BF16 shard logits, at the configured draft length (the AgentX contract). 0: (the gathered logits +# and the in-tree sampler run). Rubin TP2 b1/b2 (p5 microbench): 180.1 / 181.8 -> 115.1 / 113.9 us per verify. +SAMPLED_VERIFY_MAX_REQUESTS = 2 # capacity: fixed: the TP8 GEN batch +# A step verified by the in-tree path takes the fused first drafter inputs (the capture norms and the first pass's +# gather) up to this many requests. +FIRST_DRAFTER_INPUTS_MAX_REQUESTS = 0 # capacity: fixed: off at TP2 (0) +# A step verified by the fused kernels (dispatch.SHARD_ACCEPT / SAMPLED_ACCEPT) of up to this many requests takes the +# fused first drafter inputs (spec_worker._m3_first_drafter_inputs: the accepted tokens as they lie, the first pass's +# gather rows, the fc_norm chunk norm or the armed capture norm, then the in-tree fc). 0: (the in-tree first +# drafter inputs run, the captures stay in-tree and the fused draft loop steps aside). +FUSED_FIRST_INPUTS_MAX_REQUESTS = 2 # capacity: fixed: the TP2 GEN batch +# On a fused-verify step of up to this many verify rows the Eagle3 captures write rmsnorm(bf16(h + residual)) * +# fc_norm[i] in one kernel (ops/spec/m3_capture_norm) once the fc_norm weights are stacked; the fused first drafter +# inputs then skip their chunk norm. 0: (the in-tree capture, then the chunk norm). +CAPTURE_NORM_MAX_ROWS = 8 # capacity: fixed: the TP2 verify rows (2 requests x 4) +# The first draft pass's per-layer fc_norm over the captured columns in one kernel (ops/spec/m3_chunk_norm) at up +# to this many verify rows; wider first passes take the in-tree first drafter inputs. 0:. +CHUNK_NORM_MAX_ROWS = 8 # capacity: fixed: the TP2 verify rows (2 requests x 4) +# The fused draft loop's two all-reduce boundaries per pass (o_proj -> post-attention norm, down -> final norm) of up +# to this many rows run the multicast all-reduce + residual + RMSNorm (kernels/m3_mc_allreduce, opts bit 16: the +# drafter's stock rounding, normalizing the unrounded FP32 sum as FlashInfer fused_add_rmsnorm does) instead of the +# built-in all-reduce + RMSNorm. Kill: 0. Rubin TP2 (p16 microbench): 10.0..11.0 -> 4.6..5.3 us per boundary. +DRAFT_MC_AR_MAX_ROWS = 8 # capacity: fixed: the TP8 verify rows +# The drafter's down -> final-norm boundary (the all-reduce right before each pass's draft step) runs without the +# max-smem carveout (opts bit 8, which AR_OPTS_BY_M sets for 3..6 rows): the carveout made the 212-CTA draft step wait +# for the reduce's CTAs (TP8 census r3: pass 0 at 4 rows ran its reduce at 233,472 B smem, passes 1-2 at 16,384, and +# pass 0's step +10 us; TP4's draft-step penalty, same mechanism). Bit-identical. Kill: True. +DRAFT_MC_AR_CARVEOUT = False +# The drafter's o_proj and down projection, on the copied split-K kernel (DRAFT_LOWM_OWN_SPLITK), also push their rows +# into the drafter boundary all-reduce's exchange (the canonical words its own push would store, at the call parity), +# which then runs prepushed 2 instead of pushing x itself; the copy releases its dependents only after its own grid +# wait and the pushes (the prepushed >= 2 contract). Bitwise. Rubin TP8 world 8, residuals chained boundary to +# boundary (push-3 782332): 8.56 -> 8.08 us per boundary (1 row), 9.21 -> 8.10 (4 rows), ~ -4 us/step. Kill: False. +DRAFT_AR_PUSH = True +# On a fused-verify step of up to this many verify rows the drafter's first fc is forked onto the aux stream right +# after the last captured target layer, overlapping the remaining target layers and the head (spec_worker). +FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover +# A sampled fused verify of up to this many verify rows on the in-tree head GEMM forks the drafter's first fc after the +# head instead of beside it (spec_worker): the fc's split-K grid otherwise holds the SMs first and the head GEMM starts +# late. Kill: 0. Rubin TP2 c1 (locked nsys 768559): the head started 17.4 us after its input was ready. +FC_AFTER_HEAD_MAX_ROWS = 8 # capacity: fixed: measured fork crossover +# The forked fc's completion reaches the first draft input through a counter instead of a stream join: the fork bumps a +# device counter after the fc, a one-thread wait kernel queued after the acceptance spins on it, and the fork is joined +# after the first draft input. That input then has one producer on its own stream (a programmatic, PDL edge) instead of +# the acceptance on another graph stream: in-model the acceptance -> draft input edge cost 2.0-2.4 us (TP2 r4.4) and +# 3.2-3.8 us (TP8 tp8-full) per step. Ordering only: the same values. Kill: False (join before the draft input). +DI_FC_FLAG = True +DI_FC_FLAG_PATH = "m3_di_fc_flag" +# A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (spec_worker): per pass the +# draft input in one launch (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_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +# 0: (the in-tree draft loop runs). Rubin TP2 per draft pass (head + argmax + exchange): 122.6 / 126.0 -> +# 102.8 / 102.8 us at b1 / b2 (BF16 fragments; the coded head is slower on sm_107: DRAFT_HEAD_CODED False). +DRAFT_STEP_MAX_REQUESTS = 2 # capacity: fixed: the TP8 GEN batch +# H-R2-3: on the fused draft loop's last pass, the step's forward-tail store and next-step staging (the C++ fold tail, +# STEP_FOLD_TAIL_CPP) run in that draft_step's last CTA after its tokens (foldTail.h foldTailWarp, the standalone tail's +# code bit for bit), removing the draft_step -> fold_tail launch edge; M3SpecStepFold.fused_tail_operands then makes +# store_in_forward skip its launch. Off until its in-model pair. Kill: False. +DRAFT_STEP_FOLD_TAIL = False +DRAFT_STEP_FOLD_TAIL_PATH = "m3_draft_step_fold_tail" +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) run FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm) at the rows its TP8 table is certified for, up to this many rows; other calls keep the +# in-tree cuBLAS GEMM. Kill: 0. Rubin TP8 shards (sweep 777890, mclk locked) at 1 row: qkv 6.15 -> 3.94, +# o 4.97 -> 2.42, gate_up 8.09 -> 5.78, down 5.62 -> 4.00 us; the replicated fc 21.50 -> 17.77 (TP2). +DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the certified low-M rows +# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (kernels/m3_lowm_splitk): bit for bit the +# same output at every certified cell (spec/test_modeling_v2_m3_lowm_splitk.py); the seam for the drafter-boundary push +# epilogue and the down projection's SwiGLU prologue. Kill: False (FlashInfer's kernel). +DRAFT_LOWM_OWN_SPLITK = True +# On top of DRAFT_SILU_CPP: where the drafter MLP's down projection runs on a split-K tactic of the target's copy +# (DRAFT_LOWM_OWN_SPLITK), its activation is built in the GEMM's B load from gate | up instead of a separate launch +# (kernels/m3_draft_lowm.draft_down_swiglu): bit for bit the activation + GEMM +# (spec/test_modeling_v2_m3_lowm_splitk.py), one launch fewer per draft pass. Kill: False. OFF in t8-3b: bitwise but +# slower (SpecTail push-6b 782695, world 8, chained: 10.14 / 10.81 -> 17.38 / 66.85 us per boundary at 1 / 4 rows, +# ~ +70 us/step); +# the down projection keeps its boundary push through draft_linear and the C++ SwiGLU (DRAFT_SILU_CPP) runs. +DRAFT_SILU_FOLD = False +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(m for m in rows if m <= DRAFT_LOWM_MAX_ROWS) + for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() +} +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (kernels/triton_sparse_decode.MQ_MIN_TOKENS); the multi-query producer needs one KV head, which the TP8 heads have +# (8 query heads on 1 KV head per rank), so identity steps of at least 17 rows take it. +SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover +SPARSE_DECODE_KV_HEADS = 1 # capacity: fixed: the tp8 KV heads per rank +# The long-row top-16 selector takes engines of 129..LONG_SELECT_MAX_BLOCKS KV blocks of 128 tokens; it holds 2, 4 or +# 9 keys per thread by the row's valid-block count. At most 128 blocks, or past the maximum, the built-in selector +# runs; 0 turns the path off. Rubin TP2 (2 KV heads), 165,521 tokens, behind the scorer, ablation per layer: +# -0.9 us at 1 request and -1.2 us at 2 requests against the built-in selector. +LONG_SELECT_MIN_BLOCKS = 129 # capacity: fixed: past the built-in selector's 128 +LONG_SELECT_MAX_BLOCKS = 9216 # capacity: fixed: the select kernels' built blocks per row +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +# The long-row select releases the sparse decode after it right after its own grid wait (the decode stages q and the +# request length before its wait) up to this many query rows; past it the trigger stays at the select's exit. On Rubin +# TP2 the early-launched long-context sparse decode (SPARSE_DECODE_LC) crowds the select's CTAs off the SMs already at +# one request (4 rows): releasing at the exit is 336 us per decode step faster at 1 request and 350 at 2 (in-model, +# graph time, 165K-context mixed draws, locked memory clock), so the early release is off at every row count. +SELECT_EARLY_MAX_ROWS = 0 # capacity: fixed: off at every row count +# With the long-context decode's single merge (m3_attn_lc8.MERGE_ONCE), that op's CTAs also pull the layer's o_proj +# weight and scales into L2 before their grid dependency (cp.async.bulk.prefetch.L2, a hint: no value changes), so the +# o_proj GEMV after it finds them L2-warm (the SM107 TP4 LC's form, sparseDecodeLc4). Kill: False. +SPARSE_DECODE_LC8_PREFETCH_OPROJ = True +# With the long-context decode's single merge (m3_attn_lc8.MERGE_ONCE), that op also records each step's selected +# block ids per layer and pulls into L2 the K/V pages + scales of the blocks the NEXT sparse layer selected at the +# previous step (65-68% of a layer's distinct selected blocks recur one step later: MSA probe 778226 at TP4), so that +# layer's sparse decode reads them L2-warm. Hints only; up to N1_MAX_ROWS query rows. Kill: False. +SPARSE_DECODE_LC8_PREFETCH_NEXT_KV = True +N1_MAX_ROWS = 128 # capacity: fixed: the next-KV prefetch's row bound +# Without the early release, the long-row select releases the sparse decode once every score of its row is loaded +# (1, before the exact rank of the candidates) instead of at its exit (0). The release still follows +# the select's own grid wait, which the sparse decode's pre-wait q / length / block-table reads rely on. The decode's +# 16-CTA cluster launch and prologue then overlap the rank: Rubin TP2 in-model, locked memory clock, AgentX draws, +# 2 instances vs the r4.3 stored A: replay period -162 us fixed, -94 at a 323K-token request. +SELECT_LONG_TRIGGER_PHASE = 1 +# The long-row select runs as kernels/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per row +# instead of one 1024-thread CTA, the same selection bit for bit, with 4, 11 or 24 keys per thread +# (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices. On the Rubin TP8 decode chain (57 layers behind the +# producer stream, 1 / 2 requests) the select is -0.15 / -0.12 us per layer at 165K tokens, 0.00 / -0.01 at 323K and +# -0.04 / -0.04 at 500K against the 1024-thread kernel. Kill: False (the 1024-thread kernel). +SELECT_LONG_DSL = True +LONG_SELECT_DSL_KEY_WIDTHS = (4, 11, 24) +# The MSA index scorer (ops/attention/m3_index_score) of a sparse decode step, up to INDEX_SCORE_MAX_REQUESTS requests +# (0 turns the path off; wider steps run the in-tree CuTe DSL scorer): its CTAs per request by request count, and up +# to INDEX_SCORE_EARLY_MAX_REQUESTS requests the early issue (Q and the first K tile right after the grid wait, one +# older block per CTA prefetched into L2 before it). The split is captured with the graph, so it is one value per +# request count over every context. Rubin TP2 (2 index heads), 165,521 tokens, behind a 58 MB PDL producer, ablation +# per layer: 1 request 3.27 us at 1536 CTAs (512: 3.76, in-tree 256-CTA scorer 6.07); 2 requests 4.97 us at 768 +# (in-tree 6.71). Within 0.2 us of each context's best split from 16K to 400K tokens except 2 requests at 400K (+0.6). +# The early issue reads the request's first page id from the MSA page table before the grid wait: exact because the +# page table is written only on the host side before the replay (attention/test_modeling_v2_m3_block_table_contract.py); +# the request length read there is a hint, every load and store is bounded by the one read after the wait. +INDEX_SCORE_MAX_REQUESTS = 2 # capacity: fixed: the TP8 GEN batch +# At 1 request the split is one wave of CTAs (1272 = 212 SMs x 6 resident at this CTA's shared memory): at the +# 1-index-head geometry (57-layer chain behind the producer, scores and the selection bitwise) -0.03 / -0.27 / +0.06 us +# per layer at 165K / 323K / 500K tokens against 1536. Kill: 1536. +INDEX_SCORE_SPLIT_K: Mapping[int, int] = MappingProxyType({1: 1272, 2: 768}) +INDEX_SCORE_EARLY_MAX_REQUESTS = 2 # capacity: fixed: measured early-issue crossover + +# Synth only (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS, fractional): the fused verify's forced accepted counts are +# drawn per request, keyed by the request's KV length at the verify (its prompt plus what it has accepted), instead of +# from the server-wide counter, so a request's forced counts do not depend on what the server ran before (A/B draws +# become history-independent). Production sampling never forces counts. False: (the shared counter, the in-tree +# sequence). +FORCE_AL_PER_REQUEST = True + +# The decode MXFP8 GEMVs (kernels/mx_gemv, SM107 TP2 table), per projection (N, K): the rows each serves, on an +# activation its producer quantized (the MoE boundary for the QKV, the attention merge for the o_proj) and on a bf16 +# one (quantized in the GEMV's prologue, bit-identical to mxfp8_quantize). SM107 graph us per projection at 1..8 rows: +# 2.2..8.6 against 4.7..10.5 for the stock MXFP8 quantize + FlashInfer CuTe DSL GEMM (kernels/mx_gemv/op.py). +# False puts every projection back on the stock decode-graph GEMM. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +MX_GEMV_TP2 = True +# MX-TAC: the decode-graph MXFP8 projections the GEMVs do not serve take one fixed per-bucket decision instead of the +# AutoTuner's per-start choice (whose L2-hot profiling left c2 bimodal at +-6%): FlashInfer's CuTe DSL quantize and +# CuTe DSL GEMM. Stock's candidates per bucket are quantize {TRT-LLM, CuTe DSL} x GEMM {CUTLASS, CuTe DSL}; on SM107 +# every TP2 projection at 1..16 rows runs fastest on this pair, cold and L2-hot (graph us, cold: 4.7..10.3 against +# 5.0..11.5 with the TRT-LLM quantize and 9.7..58 with CUTLASS), and all four candidates give the same bits (quantize +# bytes and scales, GEMM outputs). A constant is rank-agreed and start-independent by construction. +# False returns these projections to the AutoTuner. +MXFP8_DECODE_QUANT_TACTIC = 0 # MXFP8QuantizeRunner.CUTE_DSL +MXFP8_DECODE_GEMM_TACTIC = 0 # FlashInferMXFP8GemmRunner.CUTE_DSL +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) +_GEMV_PROJECTIONS = ( + QKV_SPARSE, + QKV, + O_PROJ, + SHARED_GATE_UP, + SHARED_DOWN, + DENSE_GATE_UP, + DENSE_DOWN, +) +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + {nk: _ALL_ROWS for nk in _GEMV_PROJECTIONS} if MX_GEMV_TP2 else {} +) +MX_GEMV_BF16_MAX_TOKENS: Dict[Tuple[int, int], int] = ( + {nk: DECODE_MAX_TOKENS for nk in _GEMV_PROJECTIONS} if MX_GEMV_TP2 else {} +) + +Key = Hashable + + +@dataclass(frozen=True) +class RowRange: + """The shapes one fast path takes, and the catalog entry whose cells certify each of them. + + ``keys`` are the covered shapes: row counts for most paths, (query rows per request, requests) for the attention + kernels, a variant name for thresholds. ``certifies`` maps a cell's spec of ``entry`` to the keys that cell + drives (empty for a cell of another variant). + """ + + path: str + keys: FrozenSet[Key] + entry: OpWrapper + certifies: Callable[[Mapping[str, Any]], Iterable[Key]] + + def covers(self, key: Key) -> bool: + return key in self.keys + + def certified(self) -> FrozenSet[Key]: + return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + + +def _rows(lo: int, hi: int) -> FrozenSet[int]: + return frozenset(range(lo, hi + 1)) + + +def _spec_rows(field: str, **match: Any) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + """A cell's ``spec[field]`` when the cell's other fields match ``match``.""" + + def certifies(spec: Mapping[str, Any]) -> Iterable[Key]: + if any(spec.get(name) != value for name, value in match.items()): + return () + return (spec[field],) + + return certifies + + +def _const(*keys: Key) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + return lambda spec: keys + + +def _sparse_decode_producer(spec: Mapping[str, Any]) -> Iterable[Key]: + if not spec["identity"]: + return () + rows = spec["decode_query_len"] * len(spec["seq_lens"]) + mq = rows >= SPARSE_DECODE_MQ_MIN_TOKENS and SPARSE_DECODE_KV_HEADS == 1 + return ("multi_query" if mq else "per_row",) + + +def index_score_schedule(requests: int) -> Tuple[int, bool]: + """The index scorer's (CTAs per request, early issue) at this many requests (a covered request count).""" + return INDEX_SCORE_SPLIT_K[requests], requests <= INDEX_SCORE_EARLY_MAX_REQUESTS + + +def _index_score_requests(spec: Mapping[str, Any]) -> Iterable[Key]: + """A cell certifies its request count when it drives the TP8 index head (1 per rank) with that count's schedule.""" + requests = len(spec["seq_lens"]) + if spec.get("heads") != 1 or requests not in INDEX_SCORE_SPLIT_K: + return () + if (spec["split_k"], spec["early"]) != index_score_schedule(requests): + return () + return (requests,) + + +def _long_select_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The key widths a cell's rows reach: every width up to its blocks' when rows of mixed lengths are drawn.""" + blocks = spec["num_blocks"] + widest = 2 if blocks <= 2048 else 4 if blocks <= 4096 else 9 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_KEY_WIDTHS if w <= widest) + return (widest,) + + +def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The 384-thread select's key widths a cell's rows reach (every width up to its blocks' for mixed rows).""" + blocks = spec["num_blocks"] + widest = 4 if blocks <= 4 * 384 else 11 if blocks <= 11 * 384 else 24 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_DSL_KEY_WIDTHS if w <= widest) + return (widest,) + + +EMBED_FIRSTNORM = RowRange( + "embed_firstnorm", + _rows(1, EMBED_FIRSTNORM_MAX_TOKENS), + m3_embedding_firstnorm, + _spec_rows("rows"), +) +# The identity-selection Triton producer + merge (every request within 16 blocks): the per-row producer, and from +# SPARSE_DECODE_MQ_MIN_TOKENS rows the multi-query producer (one KV head), certified at the TP8 heads on SM107. +SPARSE_DECODE_PRODUCER = RowRange( + _sparse_decode_kernel.PRODUCER_PATH, + _sparse_decode_kernel.PRODUCER_KEYS, + m3_sparse_decode, + _sparse_decode_producer, +) +# The TP8 long-context sparse decode at 8 query heads over one KV head (kernels/m3_attn_lc8): attention, split-K merge +# and the MXFP8 o_proj hand-off in one launch, for every decode step with a selected table; "prewait" up to +# m3_attn_lc8.PREWAIT_MAX_ROWS query rows, "plain" above. Exact-class (in the float64 reference's band; run-to-run +# deterministic). An empty range routes every step back to the Triton producer + merge. +SPARSE_DECODE_LC8 = RowRange( + _sparse_decode_kernel.LC8_PATH, + _sparse_decode_kernel.LC8_KEYS, + m3_sparse_decode_lc8, + lambda spec: ("prewait" if spec["prewait"] else "plain",), +) +# Every token count of a sparse layer on an NVFP4 KV cache at the TP8 heads (prefill and decode): the producer takes the +# geometry, not a row range. Kill switch: frozenset() routes every sparse layer back to the native producer. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp8"}), m3_producer, _const("nvfp4_tp8")) +INDEX_SCORE = RowRange( + "m3_index_score", + _rows(1, INDEX_SCORE_MAX_REQUESTS), + m3_index_score, + _index_score_requests, +) +LONG_SELECT = RowRange( + "long_select", + frozenset(LONG_SELECT_KEY_WIDTHS) if LONG_SELECT_MAX_BLOCKS else frozenset(), + m3_select_long, + _long_select_widths, +) +LONG_SELECT_DSL = RowRange( + "long_select_dsl", + frozenset(LONG_SELECT_DSL_KEY_WIDTHS) + if LONG_SELECT_MAX_BLOCKS and SELECT_LONG_DSL + else frozenset(), + m3_select_long_dsl, + _long_select_dsl_widths, +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS) if MC_AR_TP2 else frozenset(), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=0), +) +ATTN_PUSH = RowRange( + "attn_push", + _rows(1, ATTN_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +DENSE_PUSH = RowRange( + "dense_push", + _rows(1, DENSE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +# The router of a pure generation step of 9..16 rows on the generic layer (above the MoE engine's 8; the TP2 GEN batch 4 +# with 3 drafts): the TF32 tensor-core router (minimax_m3_router_tc, 16-CTA clusters, one reduction order), the +# arithmetic the built-in gate takes on SM107 at these rows (2.9-3.5 us/call against ~9.8 us for the cuBLAS TF32 +# split-K router; closer to the fp32 product: max |dlogit| 1.7e-6 vs 6.7e-5). Devices that cannot co-schedule 16-CTA +# clusters keep the fp32 router. Context steps never take it. +ROUTER_TC_GEN_MIN_ROWS = 9 # capacity: fixed: past the MoE engine's 8 rows +ROUTER_TC_GEN_MAX_ROWS = 16 # capacity: fixed: the shared gate's router_tc range +ROUTER_TC_GEN = RowRange( + "m3_router_tc", + _rows(ROUTER_TC_GEN_MIN_ROWS, ROUTER_TC_GEN_MAX_ROWS), + m3_router_tc, + _spec_rows("m"), # every cell checks the logits, with or without the NVFP4 activation +) +MOE_PUSH = RowRange( + "moe_push", + _rows(1, MOE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MOE_ENGINE = RowRange( + "m3_moe_engine", _rows(1, MOE_ENGINE_MAX_TOKENS), m3_moe_engine, _spec_rows("tokens") +) +# The GEMM paths are keyed (projection (N, K), rows): each projection runs its own split-K / ring / MMA +# configuration, so a row count is certified per projection. +MX_GEMV = RowRange( + "m3_mx_gemv", + frozenset((nk, m) for nk, rows in MX_GEMV_ROWS.items() for m in rows), + m3_mx_gemv, + lambda spec: ((spec["nk"], spec["m"]),), +) +MX_GEMV_BF16 = RowRange( + "m3_mx_gemv_bf16", + frozenset((nk, m) for nk, hi in MX_GEMV_BF16_MAX_TOKENS.items() for m in _rows(1, hi)), + m3_mx_gemv_bf16, + lambda spec: ((spec["nk"], spec["m"]),), +) +# (draft length, requests) of a greedy verify step. +SHARD_ACCEPT = RowRange( + "m3_shard_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SHARD_ACCEPT_MAX_REQUESTS + 1)), + m3_shard_accept, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) + +# Rows of a fused-draft-loop all-reduce boundary (drafter hidden 6144) on the multicast all-reduce. +DRAFT_MC_AR = RowRange( + "m3_draft_mc_allreduce_norm", + _rows(1, DRAFT_MC_AR_MAX_ROWS), + m3_mc_allreduce_norm, + lambda spec: (spec["m"],) if spec.get("drafter") else (), +) + +# (draft length, requests) of a sampled verify step. +SAMPLED_ACCEPT = RowRange( + "m3_sampled_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SAMPLED_VERIFY_MAX_REQUESTS + 1)), + m3_sampled_accept, + lambda spec: ((SHARD_ACCEPT_DRAFT_LEN, spec["nreq"]),), +) + + +# (query rows per request, requests) of a decode step whose dense-layer attention runs on PrimTS split-KV decode. +FULLCTX_DECODE = RowRange( + "m3_fullctx_decode", + frozenset(FULLCTX_DECODE_RECIPES) if FULLCTX_DECODE_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + +# (query rows per request, requests) of a generation-only drafter pass that runs on the full-context decode. +DRAFTER_FULLCTX = RowRange( + "m3_drafter_fullctx", + frozenset(FULLCTX_DECODE_RECIPES) if DRAFTER_FULLCTX_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + + +def _verify_rows(max_rows: int) -> FrozenSet[int]: + """The verify rows of 1..8 requests at the shard-accept draft length, up to ``max_rows``.""" + return frozenset( + n * (SHARD_ACCEPT_DRAFT_LEN + 1) + for n in range(1, 9) + if n * (SHARD_ACCEPT_DRAFT_LEN + 1) <= max_rows + ) + + +# Requests of a fused-verify step whose first drafter inputs are the fused ones; certified by the chunk-norm cells +# at those requests' verify rows (the norm is the only kernel of its own; the gather rows and the fc are in-tree). +FUSED_FIRST_INPUTS = RowRange( + "m3_fused_first_inputs", + _rows(1, FUSED_FIRST_INPUTS_MAX_REQUESTS), + m3_chunk_norm, + lambda spec: ( + (spec["m"] // (SHARD_ACCEPT_DRAFT_LEN + 1),) + if spec["m"] % (SHARD_ACCEPT_DRAFT_LEN + 1) == 0 + else () + ), +) +# Verify rows of a fused-verify step whose Eagle3 captures apply the fc_norm. +CAPTURE_NORM = RowRange( + "m3_capture_norm", _verify_rows(CAPTURE_NORM_MAX_ROWS), m3_capture_norm, _spec_rows("m") +) +# Verify rows of a first draft pass whose fc_norm runs as the chunk norm. +CHUNK_NORM = RowRange( + "m3_chunk_norm", _verify_rows(CHUNK_NORM_MAX_ROWS), m3_chunk_norm, _spec_rows("m") +) + +# ((N, K), rows) of a drafter BF16 linear on its low-M kernel. +DRAFT_LOWM = RowRange( + "m3_draft_lowm", + frozenset((nk, m) for nk, rows in DRAFT_LOWM_ROWS.items() for m in rows), + m3_draft_lowm, + lambda spec: ((spec["nk"], spec["m"]),), +) +# Verify rows of a step whose target logits come from the die-placed head (in-tree GEMM halves per die). + +# (draft length, requests) of a fused draft loop, and the rows of its passes' draft input: the first pass's +# requests x (draft length + 1), one per request on the later passes. +DRAFT_STEP = RowRange( + "m3_draft_step", + frozenset((DRAFT_STEP_DRAFT_LEN, n) for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1)), + m3_draft_step, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) +DRAFT_INPUT = RowRange( + "m3_draft_input", + frozenset( + rows + for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1) + for rows in (n, n * (DRAFT_STEP_DRAFT_LEN + 1)) + ), + m3_draft_input, + _spec_rows("m"), +) + +# The dense layers' down projection takes the gate_up GEMV's output and applies the swigluoai activation in its +# quantize prologue (kernels/mx_gemv x_swiglu, bit for bit the Triton silu_and_mul kernel), instead of a separate +# silu_and_mul launch between the two GEMVs; rows the bf16 GEMV serves. Kill constant: False runs the activation +# kernel as before. +# OFF: bitwise but SLOWER (Host TP2 same-job pair 782860 / 782935: +60 / +65 us/step fixed, both orders; every GEMV CTA +# re-activates gate | up in its quantize prologue, redundantly across the N tiles). Same lesson as the drafter fold. +DENSE_SWIGLU_FOLD = False +DENSE_SWIGLU_CPP_PATH = "m3_dense_swiglu" + +FAST_PATHS: Tuple[RowRange, ...] = ( + EMBED_FIRSTNORM, + FULLCTX_DECODE, + DRAFTER_FULLCTX, + SPARSE_DECODE_PRODUCER, + SPARSE_DECODE_LC8, + PRODUCER, + INDEX_SCORE, + LONG_SELECT, + LONG_SELECT_DSL, + MC_ALLREDUCE, + ATTN_PUSH, + DENSE_PUSH, + MOE_PUSH, + MOE_ENGINE, + ROUTER_TC_GEN, + MX_GEMV, + MX_GEMV_BF16, + SHARD_ACCEPT, + SAMPLED_ACCEPT, + DRAFT_MC_AR, + DRAFT_STEP, + DRAFT_INPUT, + DRAFT_LOWM, + FUSED_FIRST_INPUTS, + CAPTURE_NORM, + CHUNK_NORM, +) + +# -------------------------------------------------------------------------------------------------------------------- +# Selection: pure functions of the step shape +# -------------------------------------------------------------------------------------------------------------------- + + +def plan_for(num_tokens: int) -> DecodePlan: + for plan in DECODE_PLANS: + if num_tokens <= plan.max_tokens: + return plan + raise AssertionError(f"no decode plan covers {num_tokens} tokens") + + +def is_decode_step(num_tokens: int) -> bool: + """Whether a pure generation step of ``num_tokens`` rows goes to the decode target.""" + return num_tokens <= DECODE_MAX_TOKENS + + +def moe_path(num_tokens: int) -> Optional[str]: + """The MoE kernel of a decode step: the engine, or None (generic).""" + if MOE_ENGINE.covers(num_tokens): + return MOE_ENGINE.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# The dense layers' swigluoai activation as this target's C++ op (ops/gemm/m3_dense_swiglu, +# minimaxM3Sm107Tp8/swigluOai.cu, LongAttn c267035d2f: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, +# beta 1), PDL-launched with an early release so the down GEMV stages its weights during it; the down GEMV reads the +# activation only after its own grid wait. Indexer kdev17 783285 (TP2 chain): -0.69 us per dense layer vs the Triton +# kernel; LongAttn gate a 783511 (oracle 23/23). False runs the Triton kernel. +DENSE_SWIGLU_CPP_PATH = "m3_dense_swiglu" +# FlashInfer's CuTe DSL MXFP8 GEMM on the producer's MXFP8 activation (rows the GEMV does not serve). +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +# The MXFP8 quantize (tuned backend) and FlashInfer's CuTe DSL GEMM: every row the decode kernels do not serve. +GENERIC_GEMM_PATH = "mxfp8_gemm" + + +def linear_path(n: int, k: int, num_tokens: int, prequant: bool, low_m_gemm: bool) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. ``low_m_gemm`` is accepted for the call sites' + symmetry: no low-M GEMM configuration is certified at TP2, so it selects nothing. + """ + if prequant: + if num_tokens in MX_GEMV_ROWS.get((n, k), ()): + return MX_GEMV_PATH + return PREQUANT_GEMM_PATH + if num_tokens <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0): + return MX_GEMV_BF16_PATH + return GENERIC_GEMM_PATH + + +# -------------------------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A covered decode shape could not take its fast path under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() + + +def set_require(require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def missed(path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if _require[0]: + raise FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " + f"{reason}" + ) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def taken(path: str, rows: Optional[int], what: str = "") -> None: + """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line + every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" + at = "" if rows is None else f" at {rows} rows" + detail = f" ({what})" if what else "" + _log_once( + logger.info, + path, + f"taken: {rows}" + (f": {what}" if what else ""), + f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", + ) + + +def long_select_covers(num_blocks: int) -> bool: + """Whether ``LONG_SELECT`` takes rows of ``num_blocks`` KV blocks (the engine's per-row block width).""" + return LONG_SELECT_MIN_BLOCKS <= num_blocks <= LONG_SELECT_MAX_BLOCKS + + +def long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: + """The (PDL launch, early trigger, trigger phase) keys of the 384-thread select a decode step can launch, compiled + before warmup; none when ``LONG_SELECT_DSL`` covers nothing.""" + if not LONG_SELECT_DSL.keys: + return () + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + pdl = get_env_enable_pdl() + keys = [(pdl, False, SELECT_LONG_TRIGGER_PHASE)] + if SELECT_EARLY_MAX_ROWS: + keys.append((pdl, True, 0)) + return tuple(keys) + + +def long_select_fallback(num_blocks: int, rows: int) -> None: + """The in-tree selector is about to run instead of the long-row select: a miss when ``LONG_SELECT`` covers + ``num_blocks`` (the kernel declined operands its cells certify; raises under require), uncovered otherwise (rows of + at most 128 blocks, or wider than ``LONG_SELECT_MAX_BLOCKS``: the in-tree selector serves them in every mode).""" + if long_select_covers(num_blocks): + missed( + LONG_SELECT.path, + rows, + f"the long-row select declined {num_blocks}-block scores it covers", + ) + else: + span = f"{LONG_SELECT_MIN_BLOCKS}..{LONG_SELECT_MAX_BLOCKS}" + uncovered( + LONG_SELECT.path, f"{num_blocks} KV blocks, outside {span}: the built-in selector" + ) + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +# Kernel modules report from inside their calls through kernels/m3_report: the same policy. +m3_report.install( + on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed +) + + +# -------------------------------------------------------------------------------------------------------------------- +# Target constants +# -------------------------------------------------------------------------------------------------------------------- + +# The generic boundary takes the low-latency multicast all-reduce + residual + RMSNorm up to this many tokens; with +# the shared + routed add folded in, up to its own maximum: the decode boundaries' certified rows. +MC_AR_MAX_TOKENS = GENERIC_MC_AR_MAX_TOKENS +MC_AR_FOLDED_MAX_TOKENS = GENERIC_MC_AR_FOLDED_MAX_TOKENS +# The shared expert folds into the routed engine from one token (dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# SMs left to the concurrent kernels while a persistent routed engine runs. +MOE_ENGINE_RESERVED_SMS = 24 # capacity: fixed: an SM split, not a size +# Id of the target's final hidden state in Eagle3's ``layers_to_capture``. +FINAL_HIDDEN_LAYER_ID = -1 # capacity: fixed: an Eagle3 layer id, not a size +# The trtllm-gen MoE runner's routing kind for MiniMax-M3 (sigmoid + correction bias, renormalized, scaled) and its +# gated activation (SwiGLU with the per-expert alpha / beta / clamp operands). +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) +INDEX_SCORE_PATH = INDEX_SCORE.path + +# -------------------------------------------------------------------------------------------------------------------- +# Streaming: also emit a response on decoding iterations 2..FIRST_CONTENT_EMIT_ITERS (RuntimeHooks), so a request's +# first content leaves one step after its first token. 1: the stock schedule. +FIRST_CONTENT_EMIT_ITERS = 2 + +# The runtime-value dump (TRTLLM_MODELING_V2_VALUE_DUMP=1): the metadata of every fast-path and entry decision a step +# takes, recorded on eager steps and captures (replay runs no Python), printed once at exit. +# -------------------------------------------------------------------------------------------------------------------- + +_VALUE_DUMP = os.environ.get("TRTLLM_MODELING_V2_VALUE_DUMP") == "1" +_VALUES: set = set() + + +def _record(*meta: Hashable) -> None: + if _VALUE_DUMP: + _VALUES.add(meta) + + +def _print_values() -> None: + if _VALUES: + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "0")) + for meta in sorted(_VALUES, key=repr): + print(f"[modeling_v2 value dump] rank {rank}: {meta}", flush=True) + + +if _VALUE_DUMP: + atexit.register(_print_values) + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for (it is 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'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs +# -------------------------------------------------------------------------------------------------------------------- + + +def _mxfp8_linear( + core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + _record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +_PINNED_RUNNERS: dict = {} + + +def _pinned_runner(cls, dtype): + runner = _PINNED_RUNNERS.get((cls, dtype)) + if runner is None: + runner = _PINNED_RUNNERS[(cls, dtype)] = cls(dtype) + return runner + + +def _pinned_mxfp8_gemm(act, act_scale, weight, weight_scale, out_dtype): + """The decode-graph MXFP8 GEMM at the target's fixed tactic (``MXFP8_DECODE_GEMM_TACTIC``), no AutoTuner: the + same choice on every rank and every start.""" + _record("pinned_mxfp8_gemm", tuple(weight.shape), act.shape[0], MXFP8_DECODE_GEMM_TACTIC) + runner = _pinned_runner(FlashInferMXFP8GemmRunner, out_dtype) + return runner.forward([act, act_scale, weight, weight_scale], tactic=MXFP8_DECODE_GEMM_TACTIC) + + +def _pinned_mxfp8_quantize(x): + """The decode-graph MXFP8 activation quantize at the target's fixed tactic (``MXFP8_DECODE_QUANT_TACTIC``).""" + _record("pinned_mxfp8_quantize", x.shape[0], MXFP8_DECODE_QUANT_TACTIC) + return _pinned_runner(MXFP8QuantizeRunner, x.dtype).forward( + [x], tactic=MXFP8_DECODE_QUANT_TACTIC + ) + + +_MX_TAC_ENGAGED: set = set() + + +def _note_mx_tac(kind: str, n: int, k: int, rows: int) -> None: + """MX-TAC's engagement line, once per (activation kind, projection, rows).""" + key = (kind, n, k, rows) + if key not in _MX_TAC_ENGAGED: + _MX_TAC_ENGAGED.add(key) + logger.info( + f"[M3 sm107 tp8] MX-TAC engaged: pinned MXFP8 decision (quantize tactic " + f"{MXFP8_DECODE_QUANT_TACTIC}, GEMM tactic {MXFP8_DECODE_GEMM_TACTIC}) " + f"for the [{n}, {k}] projection at {rows} rows, {kind} activation" + ) + + +def _gemv_variants(n: int, k: int, push: Optional[str], world: int): + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch routes to an [N, K] + projection's GEMV: an MXFP8 activation on the rows ``MX_GEMV_ROWS`` gives it, a bf16 one up to + ``MX_GEMV_BF16_MAX_TOKENS``; each pushing when the projection's boundary push covers the rows. ``push``: "attn" + (the o_proj into the attention boundary), "dense" (the dense MLP's down projection into the post-MLP boundary) + or "moe" (the shared expert's down projection into the MoE boundary).""" + push_rows = {"attn": ATTN_PUSH.keys, "dense": DENSE_PUSH.keys, "moe": MOE_PUSH.keys}.get( + push, frozenset() + ) + variants = [] + for rows in range(1, DECODE_MAX_TOKENS + 1): + for x_bf16, served in ( + (False, rows in MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + bucket = next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + variants.append((rows, x_bf16, (world, bucket, AR_PUSH_MODE))) + return variants + + +def _uncompiled_note() -> str: + key = mx_gemv.last_uncompiled + return "" if key is None else f" (its kernel {key} was not compiled before warmup)" + + +class _DeferredGemv: + """#117 (``QKV_PRODUCER_FOLD``): a sparse layer's QKV + index MXFP8 GEMV, recorded instead of launched. + + ``_linear(..., defer=True)`` returns one on the decode GEMV path of the 12-head packed projection at <= 8 rows with + an MXFP8 hand-off. The producer call site either runs the GEMV with the producer in its epilogue + (``_qkv_producer_fold``) or calls ``materialize`` (the plain GEMV) before anything reads the packed row; + ``placeholder`` (unwritten, the packed shape) stands for that row until then and is never read.""" + + def __init__(self, act, act_scale, weight, weight_scale, use_pdl): + self.placeholder = torch.empty( + act.shape[0], weight.shape[0], dtype=torch.bfloat16, device=act.device + ) + self.act, self.act_scale = act, act_scale + self.weight, self.weight_scale = weight, weight_scale + self.use_pdl = use_pdl + # Set once the record ran, fused (the producer GEMV) or plain (``materialize``). + self.done = False + + def materialize(self) -> torch.Tensor: + """Launch the plain GEMV this record stands for; returns the packed row the projection would have returned.""" + self.done = True + rows, (n, k) = self.act.shape[0], self.weight.shape + output = m3_mx_gemv( + self.act, self.act_scale, self.weight, self.weight_scale, True, None, self.use_pdl, None + ) + if output is None: + raise RuntimeError( + f"the deferred [{n}, {k}] GEMV declined at {rows} rows" + _uncompiled_note() + ) + taken(MX_GEMV_PATH, rows, f"[{n}, {k}] mxfp8 in") + return output + + +def _linear( + core, + x: torch.Tensor, + wk: str, + sk: str, + low_m_gemm: bool, + push: Optional[tuple] = None, + defer: bool = False, +) -> Tuple[Any, bool]: + """One MXFP8 projection ``x @ w[wk]^T``; returns (output, whether its GEMV pushed into ``push``). + + Outside decode CUDA graphs: the engine-tuned MXFP8 quantize + GEMM. Inside them the projection's (N, K) and the + step's rows pick one path (``linear_path``): + + 1. an activation its producer already quantized to MXFP8 (``m3_prequant.SLOT``, taken here): the tcgen05 decode + GEMV on the rows it serves, else the decode-graph MXFP8 GEMM on that activation; + 2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; + 3. otherwise the decode-graph MXFP8 quantize + GEMM. + + MX-TAC: on a decode graph whose MXFP8 projections the engine tunes (``tune_decode_graph_backends``), the GEMM + (and, on a bf16 activation, the quantize) the GEMVs do not serve take one fixed decision + (``MXFP8_DECODE_QUANT_TACTIC`` / ``MXFP8_DECODE_GEMM_TACTIC``: FlashInfer's CuTe DSL quantize and CuTe DSL GEMM) + instead of the AutoTuner's per-start choice (whose L2-hot profiling left c2 bimodal at +-6%). On SM107 every TP2 + projection at 1..16 rows runs fastest on this pair, cold and L2-hot (graph us, cold: 4.7..10.3 against 5.0..11.5 + with the TRT-LLM quantize and 9.7..58 with CUTLASS), and all four candidates give the same bits (quantize bytes + and scales, GEMM outputs). A constant is rank-agreed and start-independent by construction. + + A kernel declining rows its path covers is a miss; the chain then continues down the list. ``push`` (an + ``m3_mx_gemv`` ``ar`` tuple) is the boundary all-reduce exchange the GEMV writes its rows into. ``low_m_gemm`` + marks the projections the fused low-M MXFP8 GEMM may take (``LOW_M_GEMM_MAX_TOKENS``: no TP2 configuration is + built, so ``linear_path`` never selects it). ``defer`` (``QKV_PRODUCER_FOLD``, the sparse layers' packed QKV + projection): on path 1's GEMV at <= 8 rows of the 12-head projection, a ``_DeferredGemv`` is returned in place of + the output and nothing is launched. + """ + w, s = core.w[wk], core.w[sk] + x = x.contiguous() + mode = _mx_mode(core.mxfp8) + if not mode.decode_graph: + return _mxfp8_linear(core, x, w, s, mode), False + use_pdl = get_env_enable_pdl() + rows = x.shape[0] + n, k = w.shape + prequant = m3_prequant.SLOT.take(x) + if prequant is None and x.dtype != torch.bfloat16: + path = GENERIC_GEMM_PATH + else: + path = linear_path(n, k, rows, prequant is not None, low_m_gemm) + _record("linear", (n, k), rows, path, push is not None) + if prequant is not None: + act, act_scale = prequant + if defer and path == MX_GEMV_PATH and n == mx_gemv.PROD_N and rows <= 8: + return _DeferredGemv(act, act_scale, w, s, use_pdl), False + if path == MX_GEMV_PATH: + output = m3_mx_gemv(act, act_scale, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] mxfp8 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if path == PREQUANT_GEMM_PATH: + uncovered(MX_GEMV_PATH, f"the [{n}, {k}] projection at rows the GEMV does not serve") + if mode.tuned_quantize_gemm: + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("mxfp8", n, k, rows) + return output, False + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16(x, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] bf16 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, rows, f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note() + ) + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + if mode.tuned_quantize_gemm and x.dtype == torch.bfloat16: + act, act_scale = _pinned_mxfp8_quantize(x) + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("bf16", n, k, rows) + return output, False + return _mxfp8_linear(core, x, w, s, mode), False + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer boundaries: all-reduce + residual + RMSNorm +# -------------------------------------------------------------------------------------------------------------------- + + +def _can_mc(x, folded: bool = False) -> bool: + """Whether the generic boundary takes the multicast all-reduce + residual + RMSNorm: the rows + ``MC_AR_MAX_TOKENS`` / ``MC_AR_FOLDED_MAX_TOKENS`` (with the shared + routed add folded in) give it.""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (MC_AR_FOLDED_MAX_TOKENS if folded else MC_AR_MAX_TOKENS) + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + +def _mc_allreduce_norm(core, x, x2, residual, norm_w, *, quant=False, prepushed=0, capture=None): + """The multicast all-reduce + optional second addend + residual + RMSNorm (``m3_mc_allreduce_norm``), or None + without a workspace or for operands the kernel does not take (the same on every rank). With ``quant`` the normed + rows are also published, MXFP8-quantized, to ``m3_prequant.SLOT`` for the next MXFP8 projection; ``capture`` = + (out, fc_norm weight, eps) also writes the layer's Eagle3 capture norm of the updated residual.""" + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + if workspace is None: + return None + _record( + "m3_mc_allreduce_norm", x.shape[0], x2 is not None, quant, prepushed, capture is not None + ) + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm_w, + core.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + capture=capture, + ) + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm.""" + _record("m3_allreduce", x.shape[0]) + return m3_allreduce( + x, + core.allreduce, + AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ), + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The MoE: router, routed experts (trtllm-gen NVFP4), shared expert +# -------------------------------------------------------------------------------------------------------------------- + + +@dataclass +class _MoeScalars: + """The routed experts' kernel scalars of one layer, derived from the checkpoint's per-tensor scales once the + weights are real (``MiniMaxM3Tp8ModelingV2._derive_moe``).""" + + fc31_input_scale: torch.Tensor + fc2_input_scale: torch.Tensor + fc31_alpha: torch.Tensor + fc2_alpha: torch.Tensor + fc31_scale_c: torch.Tensor + act_alpha: torch.Tensor + act_beta: torch.Tensor + act_clamp: torch.Tensor + + +def _router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The fp32 router logits of MoE layer ``i`` when the MoE runs unfused (the decode steps' engines route + in-kernel): ``F.linear(float(h), w)``, the built-in gate's arithmetic. No router kernel serves TP2 rows + (``ROUTER_GEMV_MAX_TOKENS`` = ``ROUTER_TC_MAX_TOKENS`` = 0), so the experts quantize their bf16 input + themselves.""" + _record("m3_fp32_linear", h.shape[0] if h.dim() == 2 else 0) + return m3_fp32_linear(h, core.w[f"l{i}_router"]) + + +def _decode_router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The router logits of a pure generation step on the generic layer: ``ROUTER_TC_GEN``'s rows take the TF32 + tensor-core router when the device co-schedules 16-CTA clusters (decided at prebuild); every other step the fp32 + router (``_router_logits``).""" + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + ROUTER_TC_GEN.covers(rows) + and core.router_tc_cluster16 + and h.dtype == torch.bfloat16 + and h.stride(1) == 1 + and (h.stride(0) * 2) % 16 == 0 + ): + taken(ROUTER_TC_GEN.path, rows) + _record("m3_router_tc", rows) + return m3_router_tc(h, core.w[f"l{i}_router"], None, get_env_enable_pdl())[0] + return _router_logits(core, i, h) + + +def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): + """The trtllm-gen NVFP4 MoE runner over this rank's tensor-parallel slice of every expert, routing in-kernel. + + ``x_sf`` None means ``x`` is the bf16 activation, quantized here with the experts' global input scale. Returns + the finalized rows, or the unfinalized ``[gemm2 out, expert scales, expanded->permuted index]``. + """ + moe = core.moe[i] + if x_sf is None: + x, x_sf = fp4_quantize(x, moe.fc31_input_scale, 16, False, False) + rows = x.shape[0] + x_sf = x_sf.view(rows, -1).flatten() + _record("fp4_block_scale_moe_runner", rows, do_finalize) + outs = fp4_block_scale_moe_runner( + router_logits, + core.w[f"l{i}_router_bias"], + x, + x_sf.view(torch.float8_e4m3fn), + core.w[f"l{i}_e_w31"], + core.w[f"l{i}_e_s31"].view(torch.float8_e4m3fn), + None, + moe.act_alpha, + moe.act_beta, + moe.act_clamp, + core.w[f"l{i}_e_w2"], + core.w[f"l{i}_e_s2"].view(torch.float8_e4m3fn), + None, + moe.fc31_scale_c, + moe.fc31_alpha, + moe.fc2_alpha, + core.num_experts, + core.top_k, + None, + None, + core.moe_inter, + 0, + core.num_experts, + core.routed_scaling_factor, + ROUTING_MINIMAX2, + do_finalize, + act_type=ACT_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=core.max_num_tokens, + use_dp=False, + ) + return outs[0] if do_finalize else outs + + +def _split_chunk(rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +# -------------------------------------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------------------------------------- + + +def _nvfp4_layer(manager, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _producer_covered(core) -> bool: + """Whether ``PRODUCER`` covers this target's sparse layers: an NVFP4 KV cache at the TP8 heads (8 query, one KV + and one index head). A covered layer that cannot take ``m3_producer`` is a miss.""" + return core.main_kv_is_nvfp4 and (core.heads, core.kv_heads, core.index_heads) == (8, 1, 1) + + +def _producer_layout(kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 pools are the layout ``m3_producer`` is built for: [pages, 2, 1, 128, 64] with + its scale pool [pages, 2, 1, 128, 8].""" + return ( + kv_data is not None + and kv_scale is not None + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride()[2:] == (128 * 64, 64, 1) + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride()[2:] == (128 * 8, 8, 1) + and kv_scale.shape[0] == kv_data.shape[0] + ) + + +def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or data.shape[0] != scales.shape[0] + or tuple(data.shape[1:]) != (2, num_heads, 128, 64) + or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) + or data.dtype not in (torch.int8, torch.uint8) + or scales.dtype != torch.uint8 + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or cache.stride(1) < num_heads * cache.stride(2) + or cache.stride(0) < 2 * cache.stride(1) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + + +#: The step phases ("decode" / "prefill") whose producer engagement has been logged, once per process. +_PRODUCER_ENGAGED: set = set() + + +def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: + if covered: + missed(PRODUCER.path, rows, reason) + return None + + +def _sparse_qkv_producer( + core, + i: int, + packed, + position_ids, + kv_cache_manager, + msa_idx_k_cache, + msa_out_cache_loc, + deferred: Optional[_DeferredGemv] = None, +): + """The horizontal producer of a sparse layer: Q / index-Q (FP8) out, K / V / index-K written into the paged + caches (NVFP4 or FP8 main cache), for every step type (token-major). An NVFP4 cache at the TP8 heads takes + ``m3_producer`` (``PRODUCER``); a covered layer that cannot is a miss. Returns (q, index_q), or None for caches + neither kernel takes. + + ``kv_cache_manager``, ``msa_idx_k_cache`` (the per-layer index-K cache getter) and ``msa_out_cache_loc`` are the + calling target's ``step_args`` entries. ``deferred``: the QKV GEMV ``_linear`` recorded instead of launching + (``QKV_PRODUCER_FOLD``; ``packed`` is then its unwritten placeholder). The ``m3_producer`` branch runs it with the + producer in its epilogue (``_qkv_producer_fold``); any other branch launches it first (``materialize``). A None + return leaves it unlaunched.""" + covered = _producer_covered(core) + rows = int(packed.shape[0]) + if packed.dtype != torch.bfloat16 or position_ids is None: + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + manager = kv_cache_manager + if manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = manager.get_buffers(i, kv_layout="HND") + index_k_cache = msa_idx_k_cache(i) + out_cache_loc = msa_out_cache_loc + layer_uses_nvfp4 = _nvfp4_layer(manager, i) + scale_buffers = ( + manager.get_block_scale_buffers(i, kv_layout="HND") if layer_uses_nvfp4 else None + ) + kv_quant_scale = core.inv_kv_scales[i] if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supported_main = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_fused_write_layout(buffers, scale_buffers, core.kv_heads) + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale.is_cuda + and kv_quant_scale.dtype == torch.float32 + and kv_quant_scale.numel() >= 3 + and kv_quant_scale.is_contiguous() + ) + else: + supported_main = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, core.kv_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= core.kv_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + cos_sin = core.rotary_cos_sin + supported_rope = ( + cos_sin.is_cuda + and cos_sin.dtype == torch.float32 + and cos_sin.is_contiguous() + and cos_sin.dim() == 3 + and tuple(cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main + or not supported_index + or not supported_rope + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < rows + ): + return _producer_miss( + covered and layer_uses_nvfp4, + rows, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the layout the " + "producer is built for", + ) + w = core.w + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + layer_uses_nvfp4 + and (core.heads, core.kv_heads, core.index_heads) == (8, 1, 1) + and _producer_layout(buffers, scale_buffers) + ): + if deferred is not None: + folded = _qkv_producer_fold( + core, + i, + deferred, + buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + cos_sin, + pos, + ) + if folded is not None: + return folded + packed = deferred.materialize() + _record("m3_producer", rows) + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + core.heads, + get_env_enable_pdl(), + PRODUCER_EARLY_TRIGGER, + PRODUCER_PREWAIT_STEP_INPUTS, + ) + phase = "decode" if rows <= DECODE_MAX_TOKENS else "prefill" + if phase not in _PRODUCER_ENGAGED: + _PRODUCER_ENGAGED.add(phase) + logger.info( + f"[M3 sm107 tp8] fast path engaged: m3_producer ({phase}, T={rows}, " + f"early_trigger={PRODUCER_EARLY_TRIGGER}, prewait_step_inputs={PRODUCER_PREWAIT_STEP_INPUTS})" + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool {tuple(scale_buffers.shape)} is not the " + "producer's layout", + ) + if deferred is not None: + packed = deferred.materialize() + _record("m3_qkv_insert", rows, layer_uses_nvfp4) + q, index_q = m3_qkv_insert( + packed.contiguous(), + buffers.view(torch.uint8) if layer_uses_nvfp4 else buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.heads, + core.kv_heads, + core.index_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + ) + return q.flatten(1), index_q.flatten(1) + + +def _qkv_producer_fold( + core, + i: int, + deferred: _DeferredGemv, + buffers, + scale_buffers, + index_k_cache, + out_cache_loc, + kv_quant_scale, + cos_sin, + pos, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """#117: the deferred QKV GEMV with layer ``i``'s producer in its epilogue (``mx_gemv.m3_mx_gemv_producer``), the + same outputs and cache bytes as the GEMV + ``m3_producer``; None (nothing launched) when the kernel declines.""" + rows = int(deferred.act.shape[0]) + device = deferred.act.device + q = torch.empty(rows, core.heads, core.head_dim, dtype=torch.float8_e4m3fn, device=device) + index_q = torch.empty(rows, 1, core.head_dim, dtype=torch.float8_e4m3fn, device=device) + w = core.w + _record("m3_qkv_producer_fold", rows) + ok = mx_gemv.m3_mx_gemv_producer( + deferred.act, + deferred.act_scale, + deferred.weight, + deferred.weight_scale, + q, + index_q, + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc, + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + use_pdl=deferred.use_pdl, + ) + if not ok: + missed( + QKV_PRODUCER_FOLD_PATH, + rows, + "the producer GEMV declined the step" + _uncompiled_note(), + ) + return None + deferred.done = True + taken(QKV_PRODUCER_FOLD_PATH, rows, "[1536, 6144] mxfp8 in + producer") + return q.flatten(1), index_q.flatten(1) + + +def _qk_norm_rope(core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool) -> torch.Tensor: + """The per-head Gemma RMSNorm of Q and K and the partial NeoX RoPE of a dense layer's packed [Q | K | V], in one + kernel: in place, or an FP8 copy (Q / K normed and roped, V cast) when ``out_fp8``.""" + qkv = qkv.contiguous() + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + args = ( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + 1.0, # factor: no YARN + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + _record("qk_norm_rope", qkv.shape[0], out_fp8) + if out_fp8: + return m3_qk_norm_rope_fp8(*args) + return fused_qk_norm_rope(*args) + + +def _dense_fp8_op_applies(qkv: torch.Tensor, position_ids) -> bool: + """Whether a dense layer's FP8 q/k/v come from this target's op (``minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8``, + else the BF16 fused_qk_norm_rope output cast to FP8).""" + return ( + DENSE_QKV_FP8_VIA_BF16 + and hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8") + and position_ids is not None + and qkv.dtype == torch.bfloat16 + ) + + +def _dense_qk_norm_rope_fp8( + core, i: int, qkv: torch.Tensor, position_ids, kv_dest: Optional[tuple] = None +) -> torch.Tensor: + """A dense layer's FP8 q/k/v: QK norm + RoPE with Q/K rounded to BF16 and then to E4M3 and V cast, the bits of the + BF16 fused_qk_norm_rope output cast to FP8 (K is stored in the FP8 cache), in one launch of this target's op + (``DENSE_QKV_FP8_VIA_BF16``); without the op, the BF16 kernel and the cast. + + With ``kv_dest`` = (the layer's HND pool view, the step's int32 cache slots, live tokens, a [tokens, q heads, 128] + FP8 buffer) the op also stores the live tokens' K and V rows into the pool (the bytes ``write_msa_phase_kv`` + stores, one launch instead of about ten) and every row's Q into the buffer; ``kv_dest`` is passed only where the + op is known to run (``_dense_fp8_op_applies``).""" + if not _dense_fp8_op_applies(qkv, position_ids): + return fp8_e4m3_saturate(_qk_norm_rope(core, i, qkv, position_ids, False)) + _record("m3_dense_qk_norm_rope_fp8", qkv.shape[0], kv_dest is not None) + args = ( + qkv.contiguous(), + core.heads, + core.kv_heads, + core.kv_heads, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, # NeoX RoPE + position_ids.reshape(-1).contiguous().to(torch.int32), + True, # Gemma norms + ) + if kv_dest is not None: + return torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv(*args, *kv_dest) + return torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + + +def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: + """The dense MLP's SwiGLU-OAI activation: this target's C++ op (``m3_dense_swiglu``, minimaxM3Sm107Tp8/ + swigluOai.cu: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, beta 1, PDL-launched with an early + release so the down GEMV stages its weights during it) at the rows its cells certify, in any graph or eager step; + otherwise the Triton kernel (``m3_silu_and_mul``). Indexer kdev17 (TP2 chain): -0.69 us per dense layer vs the + Triton kernel. Every decline is logged once.""" + path = DENSE_SWIGLU_CPP_PATH + rows = gu.shape[0] + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + params = (core.swiglu_alpha, core.swiglu_limit, core.swiglu_beta) + act = None + if rows > max_rows: + uncovered(path, f"{rows} rows, past its certified {max_rows}") + elif params != (DENSE_ALPHA, DENSE_LIMIT, DENSE_BETA): + uncovered(path, f"activation parameters {params}") + else: + act = m3_dense_swiglu( + gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + ) + if act is None: + uncovered(path, "the gate_up layout") + else: + _record("m3_dense_swiglu", rows) + taken(path, rows, "dense swigluoai") + if act is None: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return act + + +def _fp8_index_qk_norm_rope( + core, i: int, idx_qk: torch.Tensor, position_ids, msa_idx_k_cache, msa_out_cache_loc +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + _record("m3_fp8_indexer_qk_norm_rope", rows) + return m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer helpers both targets' bodies call: the Eagle3 capture hook and the norm after a layer +# -------------------------------------------------------------------------------------------------------------------- + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + +def _next_norm(core, i: int) -> torch.Tensor: + """The norm the boundary after layer ``i`` applies: the next layer's input norm, the final norm after the last.""" + return core.w[f"l{i + 1}_in_norm"] if i + 1 < core.num_layers else core.w["final_norm"] + + +# -------------------------------------------------------------------------------------------------------------------- +# Kernel prebuilds the core runs before the executor's warmup +# -------------------------------------------------------------------------------------------------------------------- + + +def _in_tree_scorer(): + """The in-tree CuTe DSL index scorer's runner (importing it registers ``cute_dsl_minimax_m3_index_decode_score``); + imported on first use, as the CuTe DSL ops are.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLMiniMaxM3IndexDecodeScoreRunner, + ) + + return CuteDSLMiniMaxM3IndexDecodeScoreRunner + + +def _dummy_scorer_operands(dtype, heads, qlen, dim, page, device): + """One request of two blocks: real-shaped operands that load a compiled scorer module.""" + rows, blocks = 1, 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros(rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view(rows, blocks) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + return idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + + +#: Mixed (context + generation) steps run their generation rows on the built-in sparse decode +#: (``MsaDecodeFmha._run_sparse``). Its Triton producer and merge specialize on the split-K factor (16 / 8 / 4 / 2 / 1 +#: by the row count on SM107), on the row count's divisibility and on the 16-byte alignment of the request rows (a +#: mixed step's generation requests follow its context requests, so its block-table and length rows start mid-buffer), +#: so ``_prebuild_builtin_sparse_decode`` launches them per decode row count up to this many rows, both alignments. +BUILTIN_SPARSE_DECODE_PREBUILD_MAX_ROWS = ( # capacity: fixed: the prebuilt built-in decode rows + 64 # capacity: fixed: the mixed steps' generation rows prebuilt +) + + +def _prebuild_builtin_sparse_decode(query_lens, num_heads, num_kv_heads, device) -> list: + """Compile (and load) the built-in sparse decode's Triton producer and merge for every generation row count a mixed + step can have, so no step compiles one after the server is ready. Each launch mirrors the operands + ``MsaDecodeFmha._run_sparse`` passes for an NVFP4 layer (FP8 q, the HND pool's K / V planes, the scale pool's + planes, the head-major top-k view, BF16 output) on dummy values, on the current stream. Returns the row counts.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode, + ) + + page, head_dim, topk = 128, 128, MSA_REQUIRED_TOPK + pages = topk + pool = torch.zeros( + pages, 2, num_kv_heads, page, head_dim // 2, dtype=torch.uint8, device=device + ) + scales = torch.zeros( + pages, 2, num_kv_heads, page, head_dim // 16, dtype=torch.uint8, device=device + ) + global_scales = torch.ones(2, 4, dtype=torch.float32, device=device) + done = [] + for qlen in sorted(set(query_lens)): + for rows in range(qlen, BUILTIN_SPARSE_DECODE_PREBUILD_MAX_ROWS + 1, qlen): + requests = rows // qlen + q = torch.zeros(rows, num_heads, head_dim, dtype=torch.bfloat16, device=device).to( + torch.float8_e4m3fn + ) + indexes = ( + torch.arange(topk, dtype=torch.int32, device=device) + .repeat(rows, num_kv_heads, 1) + .contiguous() + ) + # One leading context request: its rows are skipped, so the generation rows start mid-buffer. + block_rows = torch.arange(topk, dtype=torch.int32, device=device).repeat( + requests + 1, 1 + ) + lengths = torch.full((requests + 1,), topk * page, dtype=torch.int32, device=device) + output = torch.empty(rows, num_heads, head_dim, dtype=torch.bfloat16, device=device) + for first in (0, 1): + minimax_m3_sparse_attn_decode( + q, + pool[:, 0], + pool[:, 1], + indexes.permute(1, 0, 2), + block_rows[first : first + requests], + lengths[first : first + requests], + sm_scale=head_dim**-0.5, + output=output, + decode_query_len=qlen, + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=global_scales[0, :1], + v_global_scale=global_scales[1, :1], + ) + done.append(rows) + return done + + +# -------------------------------------------------------------------------------------------------------------------- +# The core +# -------------------------------------------------------------------------------------------------------------------- + +# Engine states (compiled kernel + workspace) per (device, token count[, shared fold]), shared by every layer: a step +# runs the layers one after another on one stream. +_ENGINE_STATES: Dict[tuple, object] = {} + + +class MiniMaxM3Tp8ModelingV2(DecoderModel): + """The MiniMax-M3 text model of this target: its weights, its runtime resources, and the per-step dispatch.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_runtime(model_config) + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + self.mapping = mapping + # Whether the 9..16-row generation router takes the TF32 tensor-core kernel (set by prebuild_kernels). + self.router_tc_cluster16 = False + self.vocab_size = cfg.vocab_size + d = _weights.MODEL_WEIGHTS.dims(self) + self.num_layers = d.num_layers + self.hidden = d.hidden + self.heads = d.heads + self.kv_heads = d.kv_heads + self.head_dim = d.head_dim + self.index_heads = d.index_heads + self.index_dim = d.index_dim + self.q_size = d.heads * d.head_dim + self.kv_size = d.kv_heads * d.head_dim + self.moe_inter = d.moe_inter + self.num_experts = d.num_experts + self.top_k = int(cfg.num_experts_per_tok) + self.routed_scaling_factor = float(getattr(cfg, "routed_scaling_factor", 1.0)) + self.rms_eps = float(cfg.rms_norm_eps) + self.swiglu_alpha = float(getattr(cfg, "swiglu_alpha", 1.702)) + self.swiglu_beta = 1.0 # the (up + 1) offset of SwiGLU-OAI + self.swiglu_limit = float(getattr(cfg, "swiglu_limit", 7.0)) + self.max_num_tokens = model_config.max_num_tokens + self.moe_max_num_tokens = model_config.moe_max_num_tokens + self.sparse_layers = frozenset(_sparse_layer_ids(cfg)) + self.moe_layers = frozenset(_moe_layer_ids(cfg)) + mode = model_config.quant_config.quant_mode + self.main_kv_is_fp8 = bool(mode is not None and mode.has_fp8_kv_cache()) + self.main_kv_is_nvfp4 = bool(mode is not None and mode.has_fp4_kv_cache()) + # The sparse layers' QKV and index projections: one five-way packed projection, or the QKV projection and a + # separate BF16 index projection (fuse_qkv_index_projection=False) run on two streams. + self.fused_projection = bool( + getattr(model_config.sparse_attention_config, "fuse_qkv_index_projection", False) + ) + + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # The engine reads these two off the module tree: the MXFP8 GEMM state it tunes, and the opt-in to the + # automatic FlashInfer MXFP8 path on decode graphs. + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + + # The built-in fused all-reduce + residual + RMSNorm of the layer boundaries the multicast kernel does not take. + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + + # Streams: the shared / routed expert overlap and, when the MoE chunks its input, the chunking overlap. + self.aux_stream_moe = torch.cuda.Stream() + chunked = self.moe_max_num_tokens < model_config.max_num_tokens * mapping.dp_size + self.aux_stream_moe_chunk = torch.cuda.Stream() if chunked else None + self.moe_events = {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers} + self.moe_chunk_events = { + i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers + } + # The unfused sparse layers' index branch overlaps the QKV branch on its own stream. + self.aux_stream_attn = None if self.fused_projection else torch.cuda.Stream() + self.attn_events = ( + {} + if self.fused_projection + else {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.sparse_layers} + ) + + # Partial NeoX RoPE (rotary_dim of head_dim): the cos/sin table every attention kernel reads. + rope = RopeParams.from_config(cfg) + self.rope_theta = rope.theta + self.rotary_dim = int(rope.dim) + self.rotary_cos_sin = rope.create_rope_const_params(interleave=False)[1].reshape( + rope.max_positions, 2, -1 + ) + + # One MSA attention backend per layer (dense and sparse; the step metadata is the backend's). + quant_config = model_config.get_quant_config() + self.attn: List[MiniMaxM3MsaSparseAttention] = [] + for i in range(self.num_layers): + backend = create_attention( + model_config.attn_backend, + i, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=None, + quant_config=quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + q_scaling=1.0, + sparse_params=model_config.sparse_attention_config.to_sparse_params( + pretrained_config=cfg, layer_idx=i + ), + ) + backend.update_quant_config(quant_config) + if not isinstance(backend, MiniMaxM3MsaSparseAttention): + raise ValueError( + f"layer {i}: expected the MiniMax-M3 MSA backend, got {type(backend).__name__}" + ) + if i not in self.sparse_layers: + # The KV split trtllm-gen dense decode runs at, per request count. + backend.dense_decode_max_seq_len_by_batch = DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + self.attn.append(backend) + + # Built once the weights are real (``post_load_weights``): the KV (Q, K, V) scales, the MoE scalars, the + # decode engines and workspaces, the two targets. + self.kv_scales: Dict[int, torch.Tensor] = {} + self.inv_kv_scales: Dict[int, torch.Tensor] = {} + self.moe: Dict[int, _MoeScalars] = {} + self.mx_alpha: Optional[torch.Tensor] = None + self.moe_engines: Dict[int, Dict[Tuple[int, bool, bool], Optional[object]]] = {} + self.embed_owners: Dict[tuple, DecodeEmbedding] = {} + # Each sparse layer's selected block ids per query row at its last long-context decode + # (``SPARSE_DECODE_LC8_PREFETCH_NEXT_KV``, up to ``N1_MAX_ROWS`` query rows), read by the previous sparse + # layer's decode to prefetch this layer's KV. + self.n1_ids: Dict[int, torch.Tensor] = ( + { + i: torch.full( + (N1_MAX_ROWS, MSA_REQUIRED_TOPK), + -1, + dtype=torch.int32, + device=torch.device("cuda", torch.cuda.current_device()), + ) + for i in sorted(self.sparse_layers) + } + if SPARSE_DECODE_LC8_PREFETCH_NEXT_KV + else {} + ) + self._built_for: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def sm_scale(self, i: int) -> float: + return (self.head_dim**-0.5) / float(self.attn[i].q_scaling) + + # ---------------------------------------------------------------------------------------------- after the load + + def post_load_weights(self) -> None: + """Derive the kernels' state from the loaded weights, then build both targets and every decode kernel. + + Idempotent per set of weight storages: the loader's post-load walk reaches this module directly and through + the shell, and a second call on the same storages does nothing. + """ + identity = ( + self.w["embed"].data_ptr(), + self.w["l0_in_norm"].data_ptr(), + self.w["final_norm"].data_ptr(), + ) + if self._built_for == identity: + return + device = self.w["embed"].device + derived = nn.Module() + derived.register_buffer( + "mx_alpha", torch.ones(1, dtype=torch.float32, device=device), persistent=False + ) + for i in range(self.num_layers): + # No KV calibration in the checkpoint: unit (Q, K, V) scales and their reciprocals. + derived.register_buffer( + f"l{i}_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + derived.register_buffer( + f"l{i}_inv_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + for i in sorted(self.moe_layers): + for name, value in dataclasses.asdict(self._derive_moe(i)).items(): + derived.register_buffer(f"l{i}_{name}", value, persistent=False) + self.derived = derived + self.mx_alpha = derived.mx_alpha + self.kv_scales = {i: getattr(derived, f"l{i}_kv_scales") for i in range(self.num_layers)} + self.inv_kv_scales = { + i: getattr(derived, f"l{i}_inv_kv_scales") for i in range(self.num_layers) + } + self.moe = { + i: _MoeScalars( + **{ + f.name: getattr(derived, f"l{i}_{f.name}") + for f in dataclasses.fields(_MoeScalars) + } + ) + for i in sorted(self.moe_layers) + } + self._build_embed_owners() + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self._built_for = identity + + def cache_derived_state(self) -> None: + """The staged loaders' entry to the same build (``post_load_weights``).""" + self.post_load_weights() + + def _derive_moe(self, i: int) -> _MoeScalars: + """The routed experts' kernel scalars from the checkpoint's per-tensor scales (trtllm-gen NVFP4 layout). + + Global input scales: the reciprocal of the largest per-expert input scale (gate and up must agree). Per expert: + ``alpha = 1 / (input_scale * (1 / weight_scale_2))`` for the gate (fc31_alpha), the up half (its GEMM2-input + rescale ``fc31_scale_c = fc2_input_scale * up_alpha``) and the down projection; SwiGLU-OAI's activation + constants per expert, with beta and the clamp in the pre-dequant GEMM output domain (divided by fc31_alpha). + """ + w = self.w + isc1, isc3, isc2 = w[f"l{i}_e_isc1"].data, w[f"l{i}_e_isc3"].data, w[f"l{i}_e_isc2"].data + if not torch.allclose(isc1, isc3): + raise ValueError(f"layer {i}: the routed experts' w1 and w3 input scales differ") + fc31_input_scale = isc1.max().reciprocal() + fc2_input_scale = isc2.max().reciprocal() + ws2_1, ws2_3, ws2_2 = ( + w[f"l{i}_e_ws2_1"].data, + w[f"l{i}_e_ws2_3"].data, + w[f"l{i}_e_ws2_2"].data, + ) + fc31_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_1)) + up_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_3)) + fc2_alpha = 1.0 / (fc2_input_scale * (1.0 / ws2_2)) + fc31_scale_c = fc2_input_scale * up_alpha + n = self.num_experts + device = fc31_alpha.device + act_alpha = torch.full((n,), self.swiglu_alpha, dtype=torch.float32, device=device) + act_beta = torch.full((n,), self.swiglu_beta, dtype=torch.float32, device=device) + act_clamp = torch.full((n,), self.swiglu_limit, dtype=torch.float32, device=device) + act_beta.div_(fc31_alpha) + act_clamp.div_(fc31_alpha) + return _MoeScalars( + fc31_input_scale=fc31_input_scale, + fc2_input_scale=fc2_input_scale, + fc31_alpha=fc31_alpha, + fc2_alpha=fc2_alpha, + fc31_scale_c=fc31_scale_c, + act_alpha=act_alpha, + act_beta=act_beta, + act_clamp=act_clamp, + ) + + def _build_embed_owners(self) -> None: + """The fused embedding + first RMSNorm's caller-owned buffers, one per (rows, id dtype) of a decode step.""" + table, weight = self.w["embed"], self.w["l0_in_norm"] + self.embed_owners = { + (rows, dtype): DecodeEmbedding(table, weight, self.rms_eps, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + + def projections(self) -> List[Tuple[str, str, Optional[str]]]: + """Every MXFP8 projection in layer order: (weight key, scale key, the boundary its GEMV pushes into).""" + out = [] + for i in range(self.num_layers): + out.append((f"l{i}_qkv_w", f"l{i}_qkv_s", None)) + out.append((f"l{i}_o_w", f"l{i}_o_s", "attn")) + if i in self.moe_layers: + out.append((f"l{i}_sh_gu_w", f"l{i}_sh_gu_s", None)) + out.append((f"l{i}_sh_dn_w", f"l{i}_sh_dn_s", "moe")) + else: + out.append((f"l{i}_mlp_gu_w", f"l{i}_mlp_gu_s", None)) + out.append((f"l{i}_mlp_dn_w", f"l{i}_mlp_dn_s", "dense")) + return out + + def decode_query_lens(self) -> List[int]: + """The query rows per request a decode step can have: 1, and 1 + each draft length speculation runs.""" + lens = {1} + spec = getattr(self.model_config, "spec_config", None) + if spec is not None: + for drafts in ( + getattr(spec, "max_draft_len", None), + getattr(spec, "max_total_draft_tokens", None), + *(getattr(spec, "draft_len_schedule", None) or {}).values(), + ): + if drafts: + lens.add(1 + int(drafts)) + return sorted(lens) + + def prebuild_kernels(self) -> Dict[str, object]: + """Compile (and load) every kernel a decode forward dispatches to that would otherwise compile on its first + call: each projection's decode GEMVs (every row count, activation and boundary push), the MoE engines of + ``MOE_ENGINE``'s token counts, the index scorers at the decode query lengths (this target's and the in-tree + one wider steps take), the 384-thread long-row select, the dense layers' full-context decode plans, and the + built-in sparse decode a mixed step's generation rows run. Runs on the current (main) stream before the + executor's warmup, in the same order on every rank, then synchronizes; after it, GEMV compiles are frozen and + the forward only looks kernels up (a missing one is a miss). Returns the seconds each part took and the keys. + """ + times: Dict[str, object] = {} + # The 9..16-row generation router: decided once, from the device, never inside a CUDA-graph capture. + # The built-in gate's conditions, mirrored exactly: 128 experts, a hidden size in 1024s, 16-CTA clusters. + cfg = self.model_config.pretrained_config + self.router_tc_cluster16 = ( + int(cfg.num_local_experts) == 128 + and int(cfg.hidden_size) % 1024 == 0 + and bool(torch.ops.trtllm.minimax_m3_router_tc_cluster16(int(cfg.hidden_size))) + ) + logger.info( + f"[M3 prebuild] rank {mpi_rank()}: TF32 tensor-core router for {ROUTER_TC_GEN_MIN_ROWS}.." + f"{ROUTER_TC_GEN_MAX_ROWS}-row generation steps: " + + ( + "engaged (16-CTA clusters)" + if self.router_tc_cluster16 + else "not used (no 16-CTA clusters)" + ) + ) + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + from .kernels import m3_select_long_dsl as _select_long_dsl_kernel + from .kernels import m3_tokmap_guard + + device = self.w["embed"].device + # The one-request token-map guard of the stock generation preprocess (``PREPROCESS_TOKMAP_FIX``): every path + # that can run it (the drafter's full-context decode or its fallback, any trtllm-gen fallback), so it is + # installed whatever else is enabled. + m3_tokmap_guard.install( + torch.device("cuda", torch.cuda.current_device()), + DECODE_MAX_TOKENS, + lambda rows: taken(PREPROCESS_TOKMAP_PATH, rows, "one-request token map"), + lambda rows, reason: missed(PREPROCESS_TOKMAP_PATH, rows, reason), + ) + mx_gemv.freeze(False) + use_pdl = get_env_enable_pdl() + world = int(self.mapping.tp_size) + start = time.perf_counter() + for wk, sk, push in self.projections(): + weight, scale = self.w[wk], self.w[sk] + mx_gemv.ensure_compiled(weight, scale, use_pdl) + n, k = weight.shape + for rows, x_bf16, ar in _gemv_variants(n, k, push, world): + mx_gemv.precompile(weight, scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + # #117: the 12-head packed projection also takes its MXFP8 rows <= 8 with the producer epilogue. + if ( + QKV_PRODUCER_FOLD + and n == mx_gemv.PROD_N + and not x_bf16 + and ar is None + and rows <= 8 + ): + mx_gemv.precompile(weight, scale, rows, False, use_pdl=use_pdl, prod=True) + torch.cuda.synchronize() + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for i in sorted(self.moe_layers): + self._prebuild_engines(i) + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = self._prebuild_index_scorer(self.decode_query_lens(), device) + torch.cuda.synchronize() + times["scorer_s"] = time.perf_counter() - start + start = time.perf_counter() + times["select_keys"] = [] + if long_select_dsl_keys(): + # The 384-thread long-row select (``SELECT_LONG_DSL``): one kernel per launch key. + times["select_keys"] = _select_long_dsl_kernel.prebuild(long_select_dsl_keys(), device) + torch.cuda.synchronize() + times["select_s"] = time.perf_counter() - start + start = time.perf_counter() + times["fullctx_plans"] = 0 + if FULLCTX_DECODE.keys: + # The dense layers' full-context decode: one compiled PrimTS plan per covered (rows, requests). + times["fullctx_plans"] = _fullctx_decode_kernel.prebuild( + device, + {key: FULLCTX_DECODE_RECIPES[key] for key in FULLCTX_DECODE.keys}, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), + use_pdl, + ) + torch.cuda.synchronize() + times["fullctx_s"] = time.perf_counter() - start + start = time.perf_counter() + times["builtin_decode_rows"] = _prebuild_builtin_sparse_decode( + self.decode_query_lens(), self.heads, self.kv_heads, device + ) + torch.cuda.synchronize() + times["builtin_decode_s"] = time.perf_counter() - start + mx_gemv.freeze(True) + logger.info( + "[M3 prebuild] decode kernels built before warmup: GEMVs %.1f s, MoE engines %.1f s, index scorer " + "%.1f s (%d variants), long-row select %.1f s (%d variants), full-context decode %.1f s (%d plans), " + "built-in sparse decode %.1f s (%d row counts)" + % ( + times["gemv_s"], + times["moe_s"], + times["scorer_s"], + len(times["scorer_keys"]), + times["select_s"], + len(times["select_keys"]), + times["fullctx_s"], + times["fullctx_plans"], + times["builtin_decode_s"], + len(times["builtin_decode_rows"]), + ) + ) + return times + + # ------------------------------------------------------------------------------------------ the MoE engines + + def _expert_operands(self, i: int): + """The routed experts' operands the engines read (trtllm-gen NVFP4 SwiGLU-OAI layout), and the input scale.""" + w, m = self.w, self.moe[i] + tensors = [ + w[f"l{i}_e_w31"].data, + w[f"l{i}_e_s31"].data, + w[f"l{i}_e_w2"].data, + w[f"l{i}_e_s2"].data, + m.fc31_alpha, + m.fc31_scale_c, + m.act_alpha, + m.act_beta, + m.act_clamp, + m.fc2_alpha, + ] + return tensors, m.fc31_input_scale + + def _shared_weights(self, i: int): + w = self.w + return ( + w[f"l{i}_sh_gu_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_gu_s"].data, + w[f"l{i}_sh_dn_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_dn_s"].data, + ) + + def engine(self, i: int, num_tokens: int, shared: bool) -> Optional[object]: + """Layer ``i``'s prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in (and the shared + expert with ``shared``). None after a miss: the build failed or was never made (a forward never builds one).""" + return self._lookup( + i, self.moe_engines[i], (num_tokens, True, shared), MOE_ENGINE.path, num_tokens + ) + + def _lookup(self, i: int, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {i}'s engine" + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, what) + else: + missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + missed(path, num_tokens, "its build failed (logged)") + return engine + + def _prebuild_engines(self, i: int) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of layer ``i`` can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in. Whether an engine is kept is agreed + across the TP group (one host collective per engine): a build that fails on one rank drops that engine on + every rank, so the ranks never mix an engine (and its boundary push) with the generic MoE.""" + engines = self.moe_engines.setdefault(i, {}) + for num_tokens in sorted(MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in engines: + engines[key] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build( + i, MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + f"[M3 dispatch] {MOE_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + + def _try_build(self, i: int, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(i, num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {i}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, i: int, num_tokens: int, shared: bool) -> object: + """The 1..8-token engine (router, routed experts and the shared expert in one launch, at the TP8 width 384). + Rubin TP2 ablation chain (57 layers behind a 58.4 MB PDL producer, row = (P + MoE) - P, both orders): 31.3 / + 50.7 us/layer at 4 / 8 tokens against the built-in MoE's 42.4 / 55.4 (FlashInfer MXFP8 shared expert); bit for + bit the built-in block's shared.add_(routed). When the engine pushes its sum into the boundary all-reduce, + every CTA releases the PDL dependents right after its grid dependency wait instead of after the push + (``MOE_EARLY_AR_RELEASE``), so the all-reduce (which polls the pushed tagged words, not the grid) becomes + resident on the SMs the engine frees: scheduling only, the same bits (Rubin TP2 greedy early column, same-job + A,B + B,A: c1 -15 us).""" + tensors, scale = self._expert_operands(i) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + kwargs["shared_weights"] = self._shared_weights(i) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = M3MoeState( + device, + num_tokens=num_tokens, + num_ctas=sms - MOE_ENGINE_RESERVED_SMS, + shared=shared, + early_release=MOE_EARLY_AR_RELEASE, + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + # ------------------------------------------------------------------------------------------ the index scorers + + def _index_geometries(self) -> list: + """Every sparse layer's (index dtype, index heads per rank, index dim), deduplicated.""" + geoms = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + geom = (dtype, int(config.num_index_heads), int(config.sparse_index_dim)) + if geom not in geoms: + geoms.append(geom) + return geoms + + def _scorer_variants(self, query_lens) -> list: + """The scorer keys a decode step can take: every sparse layer's index geometry at each decode query length, + with each covered request count's schedule (``index_score_schedule``), where the kernel serves it.""" + runner = _index_score_kernel.IndexScoreRunner + keys = [] + for dtype, heads, dim in self._index_geometries(): + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=heads, + head_dim=dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for requests in sorted(INDEX_SCORE.keys): + split_k, early = index_score_schedule(requests) + key = (dtype, heads, int(qlen), dim, int(split_k), bool(early)) + if key not in keys: + keys.append(key) + return keys + + def _in_tree_scorer_variants(self, query_lens) -> list: + """The in-tree CuTe DSL scorer's keys (``CuteDSLMiniMaxM3IndexDecodeScoreRunner``: index dtype, index heads, + decode query length, index dim, its 256-CTA split) a decode step outside ``INDEX_SCORE`` can take: every query + length at which a step of more than ``INDEX_SCORE_MAX_REQUESTS`` requests fits in ``DECODE_MAX_TOKENS`` (e.g. + 3..8 requests decoding one row each), where that kernel serves the geometry.""" + in_tree = _in_tree_scorer() + widest = DECODE_MAX_TOKENS // (INDEX_SCORE_MAX_REQUESTS + 1) + keys = [] + for dtype, heads, dim in self._index_geometries(): + for qlen in sorted(set(query_lens)): + if qlen > widest or not in_tree.is_supported( + q_dtype=dtype, + num_heads=heads, + head_dim=dim, + page_size=in_tree.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + key = (dtype, heads, qlen, dim, in_tree.SPLIT_K) + if key not in keys: + keys.append(key) + return keys + + def _prebuild_index_scorer(self, query_lens, device) -> list: + """Compile every key of ``_scorer_variants`` and ``_in_tree_scorer_variants`` and launch each once on + real-shaped dummy operands (loading its module), on the current stream, so no decode step compiles a scorer. + The kernels' batch, token, page and block counts are symbolic: their compiles need no step metadata. Returns + the keys.""" + runner = _index_score_kernel.IndexScoreRunner + keys = self._scorer_variants(query_lens) + for key in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + q, k, bt, score, seq = _dummy_scorer_operands( + dtype, heads, qlen, dim, runner.SUPPORTED_PAGE_SIZE, device + ) + runner.kernel_cache[key](q, k, bt, score, seq) + in_tree = _in_tree_scorer() + in_tree_keys = self._in_tree_scorer_variants(query_lens) + for key in in_tree_keys: + dtype, heads, qlen, dim, split_k = key + in_tree._compile(dtype, heads, qlen, dim, split_k) + q, k, bt, score, seq = _dummy_scorer_operands( + dtype, heads, qlen, dim, in_tree.SUPPORTED_PAGE_SIZE, device + ) + in_tree.kernel_cache[key](q, k, bt, score, seq) + return keys + in_tree_keys + + # --------------------------------------------------------------------------------------------- the dispatch + + def _select_target(self, attn_metadata: AttentionMetadata): + """Which target runs this step: one host-side integer, fixed per CUDA-graph capture. A step with context rows + (a mixed step included) is prefill's; every pure generation step is decode's, whatever its width.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + capture: Callable[..., None] = _plain_capture, + capture_target: Optional[Callable[..., Optional[tuple]]] = None, + **kwargs, + ) -> torch.Tensor: + """``capture``: the Eagle3 hidden-state capture of a layer (the worker's, with its capture norm), and + ``capture_target``: where a decode boundary all-reduce writes a capture layer's capture norm itself + (``CAPTURE_NORM_IN_BOUNDARY``; the worker's, None without one).""" + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if kwargs.get("lora_params") is not None: + raise ValueError("this MiniMax-M3 target serves no LoRA adapters") + return self._select_target(attn_metadata).forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The targets +# -------------------------------------------------------------------------------------------------------------------- + + +class PrefillTarget: + """A step with context rows: prefill chunks, and the mixed steps that carry generation rows beside them. Every + layer runs in its generic form; the attention core is the MSA backend's for every row.""" + + def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + Counts and per-step host scalars are read here; device tensors are the metadata's persistent buffers, held + by reference; ``msa_idx_k_cache`` is the metadata's per-layer index-K cache getter. ``metadata`` is the + prepared object itself: the MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_tokens=md.num_tokens, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + """``capture_target`` is the decode boundaries' (unused: no prefill boundary writes a capture norm).""" + core = self.core + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) prefill CUDA graphs" + ) + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + _record("prefill", int(step["num_tokens"]), int(step["num_contexts"])) + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + hidden, residual = self._layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + return hidden + + def _swiglu_mlp( + self, h: torch.Tensor, prefix: str, low_m_gemm: bool, dense: bool = False + ) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI (the dense MLP's on this target's + C++ op, ``_dense_swiglu``), down, no reduction.""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + out, _ = _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm) + return out + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: + the MoE scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: the router logits and the experts quantizing their bf16 input + (chunked).""" + return self._routed_chunks(i, h, _router_logits(self.core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool): + """The MoE block of layer ``i`` without its reduction. Returns ``(out, routed)``: with ``return_parts`` the + shared and routed parts (routed None when the engine folded the shared expert in), else their sum in + ``out``.""" + core = self.core + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, False) + taken(MOE_ENGINE.path, rows) + folded = m3_moe_engine(h.contiguous(), engine, None, None) + return folded, None + routed_out, shared_out = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the + # same collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step): + """Attention of layer ``i`` up to its o_proj (no reduction). Returns (o_proj output, pushed): a prefill step + arms no boundary push, so pushed is always False.""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v for an FP8 cache, on this target's op; else in + # place in bf16), the MSA backend (it writes the layer's cache), o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = ( + _dense_qk_norm_rope_fp8(core, i, qkv, position_ids) + if core.main_kv_is_fp8 + else _qk_norm_rope(core, i, qkv, position_ids, False) + ) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection + and its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and + the FP8 indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the + o_proj. A layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): the MSA + backend (``m3_msa_attention``: the KV write unless a producer wrote the caches, the sparse layers' top-k block + selection, the FMHA over the context rows and any generation rows beside them).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + sparse = i in core.sparse_layers + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + +class DecodeTarget: + """A pure generation step (no context rows), of any width. + + Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids (up to 2 requests verifying 1 + 3 rows, the GEN worker's + batch), per layer: attention (o_proj all-reduce deferred, its rows pushed into the boundary all-reduce) -> + attention boundary (all-reduce + residual + post-attention RMSNorm) -> MoE (the 1..8-row engine pushing its sum + into the boundary all-reduce) or dense MLP (its down projection pushing into the post-MLP boundary) -> MoE / dense + boundary (all-reduce + shared/routed add + residual + the next layer's RMSNorm, emitting the next QKV's MXFP8 + activation). An Eagle3 capture layer runs the same decode layer and captures its feed-forward boundary's reduced + residual sum (the boundary writing the capture norm itself when the worker gives it a target). The first layer + runs the generic layer when the fused embedding does not take the step. A wider step (or one carrying embeddings) + runs every layer in its generic form, with this target's attention (the dense layers' full-context decode, the + decode selection and sparse decode where the producer wrote the caches). + """ + + def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # Set by a feed-forward boundary that wrote the layer's Eagle3 capture (see _layer's ``capture``). + self._capture_written = False + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + States the two context counts instead of reading them back: this target runs only when ``num_contexts`` is + zero, so reading them would ask the metadata to confirm what the routing already decided. + + Graph classes: counts, the generation span, the identity-selection flag, the worst-case score extent and the + cache length bound are host values fixed per capture (prepare() stages them before the forward); the block + table, lengths, cache slots, identity table and valid-block counts are the metadata's persistent buffers, held + by reference and refreshed in place every step; ``msa_idx_k_cache`` / ``msa_proxy_max_score_view`` are the + metadata's getters of the per-layer index-K cache and of the score scratch view. ``metadata`` is the prepared + object itself: the MSA backend's step (``m3_msa_attention``) and the full-context decode's K/V write + (``write_msa_phase_kv``) take it as their argument. + """ + return dict( + num_contexts=0, + num_ctx_tokens=0, + num_tokens=md.num_tokens, + num_seqs=md.num_seqs, + max_seq_len=md.max_seq_len, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + msa_decode_span=md.msa_decode_span, + msa_block_table=md.msa_block_table, + msa_seq_lens_cuda=md.msa_seq_lens_cuda, + msa_short_ok=md._msa_short_ok, + msa_short_table=md.msa_short_table, + msa_proxy_max_score_view=md.msa_proxy_max_score_view, + msa_worst_case_max_k_tiles=md.msa_worst_case_max_k_tiles, + msa_n_valid_blocks=md.msa_n_valid_blocks, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + # The decode kernels' rows (a per-capture constant): up to DECODE_MAX_TOKENS tokens of 1-D input ids. + kernels = ( + inputs_embeds is None and input_ids.ndim == 1 and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + _record("decode", num_tokens, plan.name) + hidden, residual = self._embed(input_ids, num_tokens) + else: + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + _record("decode", int(step["num_tokens"]), "generic") + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + captured = spec_metadata is not None and spec_metadata.is_layer_capture(i) + if kernels and residual is not None and captured: + # The capture layer on the decode layer (its boundaries push); the capture reads the MoE / dense + # boundary's reduced residual sum, as the generic layer's fused boundary capture does. Given a target + # (``CAPTURE_NORM_IN_BOUNDARY``), the boundary all-reduce writes the capture norm itself. + target = ( + capture_target(spec_metadata, i, num_tokens) + if capture_target is not None + else None + ) + self._capture_written = False + hidden, residual = self._layer( + i, plan, position_ids, hidden, residual, step, capture=target + ) + taken("capture_layer_fast", num_tokens, f"layer {i}") + if target is not None and self._capture_written: + taken("capture_norm_in_boundary", num_tokens, f"layer {i}") + capture(spec_metadata, i, residual, None, written=True) + else: + capture(spec_metadata, i, residual, None) + elif not kernels or residual is None or captured: + hidden, residual = self._generic_layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + else: + hidden, residual = self._layer(i, plan, position_ids, hidden, residual, step) + return hidden + + def _embed(self, input_ids: torch.Tensor, num_tokens: int): + """(hidden, residual) of the fused embedding + first RMSNorm; (embedding, None) when it does not run, so the + first layer runs its own input norm. SM107: 1.55..1.65 us against 3.4..6.0 us for the stock embedding + + RMSNorm at 1..8 rows (the gather bitwise, the norm one ulp off on ~5e-6 of the elements: summation order).""" + core = self.core + if not EMBED_FIRSTNORM.covers(num_tokens): + uncovered(EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens") + return m3_embedding(input_ids, core.w["embed"]), None + owner = core.embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + missed( + EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 after the load)" + if owner is None + else "the input ids are not contiguous", + ) + return m3_embedding(input_ids, core.w["embed"]), None + _record("m3_embedding_firstnorm", num_tokens, input_ids.dtype) + taken(EMBED_FIRSTNORM.path, num_tokens, str(input_ids.dtype)) + return m3_embedding_firstnorm( + input_ids, owner.table, owner.weight, owner.eps, owner.raw, owner.output, owner.use_pdl + ) + + def _bucket_flags(self, num_tokens: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(self.core.mapping, self.core.hidden) + workspace = None if workspace is None else m3_mc_allreduce._bucket(workspace, num_tokens) + flags = None if workspace is None else workspace.flags_for(8) + return None if flags is None else (workspace, flags) + + @staticmethod + def _no_bucket(path: str, num_tokens: int) -> None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, num_tokens: int): + """The push tuple into the attention boundary's all-reduce (the o_proj's; also the dense down projection's into + the post-MLP boundary, the same exchange), or None (after a miss). + + None comes only from the token count or from workspace state created collectively in the same eager call on + every rank, so all ranks agree. (A lone non-pushing rank would still be safe: its all-reduce pushes the same + bf16 rows into the same set A.)""" + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(ATTN_PUSH.path, num_tokens) + return None + workspace, flags = found + return ( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + AR_PUSH_MODE, + ) + + def _moe_push_targets(self, num_tokens: int): + """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None + (after a miss). + + This decision must match across ranks: prepushed 3 reads set B of every rank, which a non-pushing rank never + fills. None comes only from the token count or from state created in the same eager call on every rank + (workspace, flags, the descriptor cache), and the engine it is armed for is agreed at build time + (``_prebuild_engines``), so it does.""" + hidden = self.core.hidden + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(MOE_PUSH.path, num_tokens) + return None + workspace, flags = found + key = (workspace.mc.data_ptr(), flags.data_ptr()) + desc = self._moe_push_descs.get(key) + if desc is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor") + return None + desc = self._moe_push_descs[key] = ar_push_desc( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + flags.device, + ) + set_b = workspace.mc.data_ptr() + 2 * 4 * workspace.max_tokens * workspace.world_size * ( + hidden // 2 + ) + return desc, ( + set_b, + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + 2, + ) + + def _boundary(self, use_mc: bool, x, residual, norm_w, x2=None, quant=False, capture=None): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce. + ``capture``: the Eagle3 capture the multicast kernel also writes (``_capture_written`` records it).""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant, capture=capture) + if fused is not None: + if capture is not None: + self._capture_written = True + taken( + MC_ALLREDUCE.path, + x.shape[0], + ("x + x2" if x2 is not None else "x") + (", MXFP8 copy" if quant else ""), + ) + return fused + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + missed( + MC_ALLREDUCE.path, + x.shape[0], + m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + ) + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _prepushed( + self, i: int, x, residual, norm_w, quant, prepushed, source, x2=None, capture=None + ): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed, capture=capture + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + if capture is not None: + self._capture_written = True + taken(MC_ALLREDUCE.path, x.shape[0], f"prepushed {prepushed}, {source}") + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step, capture=None): + """Layer ``i`` on the decode kernels. ``capture``: (out, fc_norm weight, eps) for the feed-forward boundary to + write the layer's Eagle3 capture norm (``CAPTURE_NORM_IN_BOUNDARY``); a boundary that writes it sets + ``_capture_written``.""" + core = self.core + num_tokens = hidden.shape[0] + # Padded rows: the boundary all-reduce polls and re-arms rows t < M = residual.shape[0], CUDA-graph padding + # rows included. Every producer pushes all rows of the same tensor: the o_proj / shared down / dense down GEMV + # t < m = x.shape[0] (its row count comes from the host shape, not a live count), the MoE engine every row of + # the num_tokens it was built for (= num_tokens here). The all-reduce checks x.sizes() == residual.sizes(), so + # no padded row is left with a stale flag. + # The GEMV dispatch (and with it every push into a boundary all-reduce) is a decode-graph one. + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + uncovered( + "decode_graph", + "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + ) + push = plan.attn_mc_ar and decode_graph and ATTN_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if push else None + hidden, attn_pushed = self._attention(i, position_ids, hidden, step, armed) + post_norm = core.w[f"l{i}_post_norm"] + if armed is not None and attn_pushed: + hidden, residual = self._prepushed( + i, hidden, residual, post_norm, False, AR_PUSH_MODE, "o_proj push" + ) + else: + if armed is not None: + missed( + ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden, residual = self._boundary(plan.attn_mc_ar, hidden, residual, post_norm) + + norm_w = _next_norm(core, i) + if i not in core.moe_layers: + # A dense layer: its MLP's down projection pushes into the post-MLP boundary's all-reduce, as the o_proj + # does into the attention one (the same exchange; the gate_up projection never pushes). SM107 TP2: the + # two non-capture dense layers' post-MLP hops end 5.5-5.8 us after the down GEMV unpushed vs 2.4-2.5 us + # pushed; bitwise equal to the unpushed boundary (the same bf16 words). + dense_push = plan.attn_mc_ar and decode_graph and DENSE_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if dense_push else None + out, pushed = self._swiglu_mlp(hidden, f"l{i}_mlp", True, armed, dense=True) + if armed is not None and pushed: + return self._prepushed( + i, + out, + residual, + norm_w, + False, + AR_PUSH_MODE, + "dense down push", + capture=capture, + ) + if armed is not None: + missed( + DENSE_PUSH.path, + num_tokens, + "the dense MLP's down projection did not push its rows into the armed boundary all-reduce", + ) + return self._boundary(plan.attn_mc_ar, out, residual, norm_w, capture=capture) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + moe_push = None + if moe_path(num_tokens) == MOE_ENGINE.path: + # The engine runs inside the MoE; looked up (or found missing) here first so a push is armed only for an + # engine that will take it. + engine = core.engine(i, num_tokens, shared=True) + if ( + engine is not None + and plan.moe_mc_ar + and decode_graph + and MOE_PUSH.covers(num_tokens) + ): + moe_push = self._moe_push_targets(num_tokens) + shared, routed, engine_pushed, shared_pushed = self._moe( + i, hidden, return_parts=True, push=moe_push + ) + if moe_push is not None: + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + if engine_pushed: + return self._prepushed( + i, shared, residual, norm_w, quant, 2, "MoE engine push", capture=capture + ) + else: + if engine_pushed and shared_pushed: + return self._prepushed( + i, + shared, + residual, + norm_w, + quant, + 3, + "MoE engine + shared push", + x2=routed, + capture=capture, + ) + if engine_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + missed( + MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary( + plan.moe_mc_ar, shared, residual, norm_w, x2=routed, quant=quant, capture=capture + ) + + def _swiglu_mlp( + self, + h: torch.Tensor, + prefix: str, + low_m_gemm: bool, + push: Optional[tuple] = None, + dense: bool = False, + ): + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI (the dense MLP's on this target's + C++ op, ``_dense_swiglu``), down, no reduction. Returns (output, whether the down projection's GEMV pushed into + ``push``).""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm, push) + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: the MoE + scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: the router logits (the TF32 tensor-core router on ``ROUTER_TC_GEN``'s + rows, else the fp32 router) and the experts quantizing their bf16 input (chunked).""" + return self._routed_chunks(i, h, _decode_router_logits(self.core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool, push=None): + """The MoE block of layer ``i`` without its reduction. + + ``push`` = (the engine's ``arp`` descriptor, the shared down projection's ``ar`` tuple) the decode layer armed + for the MoE boundary's exchange sets, or None. Returns ``(out, routed, engine_pushed, shared_pushed)``: with + ``return_parts`` the shared and routed parts (routed None when the engine folded the shared expert in), else + their sum in ``out``. + """ + core = self.core + engine_arp, shared_ar = push if push is not None else (None, None) + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, engine_arp is not None) + taken(MOE_ENGINE.path, rows) + folded = m3_moe_engine(h.contiguous(), engine, None, engine_arp) + return folded, None, engine_arp is not None, False + routed_out, (shared_out, shared_pushed) = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False, shared_ar), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the same + # collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out, False, shared_pushed + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None, False, shared_pushed + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _generic_layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed, _, _ = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = None): + """Attention of layer ``i`` up to its o_proj (no reduction). ``push`` is the attention boundary's exchange the + o_proj may write into. Returns (o_proj output, pushed).""" + core = self.core + if i not in core.sparse_layers: + return self._dense_attention(i, position_ids, hidden, step, push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + # #117: the QKV GEMV may be deferred to the producer call site (``QKV_PRODUCER_FOLD``). + packed, _ = _linear( + core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True, defer=QKV_PRODUCER_FOLD + ) + deferred = None + if isinstance(packed, _DeferredGemv): + deferred, packed = packed, packed.placeholder + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + deferred, + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _dense_attention(self, i: int, position_ids, hidden, step, push): + """Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v: an FP8 cache, or the FP8 pool the dense layers sit + in within a hybrid NVFP4 cache on a decode step), the attention core, o_proj. + + On a covered pure generation step (``_fullctx_plan``) the attention is this target's PrimTS split-KV decode + over the layer's full causal context, and the norm/RoPE op also stores the step's new K/V rows into the + layer's paged pool and the query compact, as the decode reads it (``DENSE_KV_INSERT_FUSED``; the bytes + ``write_msa_phase_kv`` stores). Otherwise the MSA backend.""" + core = self.core + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + plan = self._fullctx_plan(i, qkv.device, step) + kv_dest = None + if ( + plan is not None + and fp8 + and DENSE_KV_INSERT_FUSED + and _dense_fp8_op_applies(qkv, position_ids) + and hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + ): + kv_dest = ( + step["kv_cache_manager"].get_buffers(i, kv_layout="HND"), + step["msa_out_cache_loc"], + plan[1], + torch.empty( + (qkv.shape[0], core.heads, core.head_dim), + dtype=torch.float8_e4m3fn, + device=qkv.device, + ), + ) + if fp8: + fused = _dense_qk_norm_rope_fp8(core, i, qkv, position_ids, kv_dest) + if kv_dest is not None: + taken(DENSE_KV_INSERT_PATH, kv_dest[2]) + else: + fused = _qk_norm_rope(core, i, qkv, position_ids, False) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + if kv_dest is not None: + q = kv_dest[3].view(q.shape[0], -1) + o = None + if plan is not None: + o = self._fullctx_decode(i, q, k, v, step, plan, kv_written=kv_dest is not None) + if o is None: + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _fullctx_plan(self, i: int, device: torch.device, step): + """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs + on ``kernels/m3_fullctx_decode`` (``FULLCTX_DECODE``): a covered (query rows per request, requests) shape over + the layer's FP8 pool. None otherwise (the MSA backend runs). + + Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 (29.3), (1, 1) + 15.7 (28.0), (4, 2) 24.5 (29.3), (1, 2) 20.2 (28.1); the stock split is capped at 16 cluster CTAs per (request, + KV head), i.e. 32 of 212 CTAs at one request. With the BF16 split-KV merge (``FULLCTX_DECODE_OWN_REDUCER``): + (1, 1) 14.5 vs 18.5, (4, 1) 17.3 vs 21.3, (1, 2) 18.5 vs 23.1, (4, 2) 23.3 vs 27.5.""" + path = FULLCTX_DECODE.path + span = step["msa_decode_span"] + if span is None or span.row_first != 0: + uncovered(path, "no decode span starting at row 0") + return None + nreq = int(step["num_seqs"]) + qlen = int(span.query_len) + if not FULLCTX_DECODE.covers((qlen, nreq)): + uncovered(path, f"{qlen} query rows x {nreq} requests is not a covered shape") + return None + if int(step["max_seq_len"]) > FULLCTX_DECODE_MAX_KV_LEN: + uncovered(path, "the cache length bound is beyond the plans' K/V capacity") + return None + num_tokens = nreq * qlen + pool, _ = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + if pool.dtype != torch.float8_e4m3fn: + missed(path, num_tokens, "the dense layer's KV pool is not FP8") + return None + from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + + own_reducer = FULLCTX_DECODE_OWN_REDUCER + if own_reducer and not _fullctx_decode_kernel.has_own_reducer(): + missed(path, num_tokens, "the BF16 split-KV merge op is not built") + return None + # Compiled before warmup (``prebuild_kernels``); the forward never compiles one. + plan = _fullctx_decode_kernel.plan_for(device, qlen, nreq, own_reducer) + if plan is None: + missed(path, num_tokens, "the PrimTS decode plan is not prebuilt") + return None + return plan, num_tokens + + def _fullctx_decode(self, i: int, q, k, v, step, plan_tokens, *, kv_written: bool): + """Dense layer ``i``'s decode attention on the PrimTS plan of ``_fullctx_plan``: the step's new K/V go into the + paged cache as the MSA backend writes them (``write_msa_phase_kv``; already there when ``kv_written``), then + every query row attends its request's full causal context.""" + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + write_msa_phase_kv, + ) + + from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + + core = self.core + plan, num_tokens = plan_tokens + nreq = int(step["num_seqs"]) + if not kv_written: + write_msa_phase_kv( + core.attn[i], + k[:num_tokens], + v[:num_tokens], + step["metadata"], + AttentionInputType.mixed, + token_offset=0, + ) + pool, subpages = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + k_pages, v_pages = _fullctx_decode_kernel.kv_page_views(pool, subpages) + # The merge also writes the o_proj's MXFP8 hand-off (``FULLCTX_DECODE_MXFP8``) where its TP8 copy runs. + mxfp8 = None + if ( + FULLCTX_DECODE_MXFP8 + and output.shape[0] == num_tokens + and _fullctx_decode_kernel.emits_mxfp8() + ): + merge_mxfp8 = _sparse_decode_kernel.MERGE_MXFP8 + if _sparse_decode_kernel.MERGE_MXFP8_MIN_TOKENS <= num_tokens <= merge_mxfp8.MAX_ROWS: + mxfp8 = merge_mxfp8.buffers(output.shape[1], output.device) + _record("m3_fullctx_decode", num_tokens, nreq, kv_written, mxfp8 is not None) + emitted = plan.run( + fp8_e4m3_saturate(q[:num_tokens].reshape(num_tokens, core.heads, core.head_dim)), + k_pages, + v_pages, + step["msa_block_table"][:nreq], + step["msa_seq_lens_cuda"][:nreq], + core.sm_scale(i), + output[:num_tokens], + mxfp8=mxfp8, + ) + if emitted: + m3_prequant.SLOT.publish(output, mxfp8[0][:num_tokens], mxfp8[1]) + taken(FULLCTX_DECODE_MXFP8_PATH, num_tokens) + taken(FULLCTX_DECODE.path, num_tokens) + return output + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step, push): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection and + its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and the FP8 + indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the o_proj. A + layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): a sparse + layer whose caches the producer wrote runs this target's selection and sparse decode when the generation span + starts at row 0; everything else the MSA backend (``m3_msa_attention``: the KV write unless a producer wrote the + caches, the sparse layers' top-k block selection, the FMHA).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + span = step["msa_decode_span"] + sparse = i in core.sparse_layers + if sparse and k_live is None and span is not None and span.row_first == 0: + self._decode_sparse_core(i, q_live, idx_q_live, step, out_live) + return output + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + def _index_scores(self, i: int, idx_q, step, num_tokens: int, qlen: int): + """The per-index-head block maxima of a non-identity decode step: this target's CuTe DSL scorer at the request + counts ``INDEX_SCORE`` covers (each count's CTAs per request and early issue, ``index_score_schedule``), the + in-tree CuTe DSL scorer beyond them. Rubin TP2 (2 index heads), 165,521 tokens, behind a 58 MB PDL producer, + per layer: 1 request 3.27 us at 1536 CTAs (in-tree 256-CTA scorer 6.07); 2 requests 4.97 us at 768 (in-tree + 6.71).""" + config = self.core.attn[i].m3_config + rows = int(step["num_seqs"]) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = step["msa_idx_k_cache"](i) + max_score = step["msa_proxy_max_score_view"]( + config.num_index_heads, step["msa_worst_case_max_k_tiles"], num_tokens + ) + if INDEX_SCORE.covers(rows): + split_k, early = index_score_schedule(rows) + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(split_k), + bool(early), + ) + if key not in _index_score_kernel.IndexScoreRunner.kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + missed( + INDEX_SCORE_PATH, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score", key) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + split_k, + early, + ) + taken(INDEX_SCORE_PATH, rows) + else: + uncovered( + INDEX_SCORE_PATH, + f"{rows} requests, past {INDEX_SCORE_MAX_REQUESTS}: the in-tree CuTe DSL scorer", + ) + _in_tree_scorer() # registers the op + _record("cute_dsl_minimax_m3_index_decode_score", rows, qlen) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + ) + return max_score + + def _select(self, i: int, max_score, step, num_tokens: int): + """The top-16 table of a non-identity decode step from the block maxima reduced to the KV head: the long-row + select (the 384-thread kernel, ``SELECT_LONG_DSL``, when compiled) on rows of ``LONG_SELECT_MIN_BLOCKS.. + LONG_SELECT_MAX_BLOCKS`` blocks, else the built-in selector (``m3_select_blocks``).""" + config = self.core.attn[i].m3_config + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + group = config.num_index_heads // config.num_kv_heads + max_score_kv = ( + max_score.view(config.num_kv_heads, group, max_score.shape[1], max_score.shape[2]).amax( + dim=1 + ) + if group > 1 + else max_score + ) + from .kernels import m3_select_long_dsl as _select_long_dsl_kernel + + num_blocks = int(max_score_kv.shape[1]) + table = None + dsl = False + if long_select_covers(num_blocks): + early_trigger = num_tokens <= SELECT_EARLY_MAX_ROWS + if LONG_SELECT_DSL.keys: + # Compiled before warmup (prebuild_kernels); a missing kernel falls back to the 1024-thread one. + key = _select_long_dsl_kernel.kernel_key( + get_env_enable_pdl(), early_trigger, SELECT_LONG_TRIGGER_PHASE + ) + dsl = key in _select_long_dsl_kernel.kernel_cache + if not dsl: + missed( + LONG_SELECT_DSL.path, + num_tokens, + f"the 384-thread select {key} was not compiled before warmup", + ) + _record("long_select", num_blocks, num_tokens, dsl) + table = long_select_blocks( + max_score_kv, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + use_pdl=get_env_enable_pdl(), + early_trigger=early_trigger, + trigger_phase=SELECT_LONG_TRIGGER_PHASE, + dsl=dsl, + ) + if table is not None: + taken((LONG_SELECT_DSL if dsl else LONG_SELECT).path, num_tokens) + else: + long_select_fallback(num_blocks, num_tokens) + _record("m3_select_blocks", num_blocks, num_tokens) + table = m3_select_blocks( + max_score_kv, + MSA_REQUIRED_TOPK, + n_valid_blocks, + config.init_blocks, + config.local_blocks, + True, + ) + return table + + def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step of a sparse layer whose caches + the producer already wrote. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise the + index scorer writes the per-index-head block maxima, reduced to the KV head, and the long-row selector takes the + top 16 (the built-in selector beyond its shapes). The long-context decode (``m3_sparse_decode`` on + ``kernels/m3_attn_lc8``) prefetches into L2 this layer's o_proj weights (``SPARSE_DECODE_LC8_PREFETCH_OPROJ``) + and the next sparse layer's K/V pages of the blocks that layer selected one step back + (``SPARSE_DECODE_LC8_PREFETCH_NEXT_KV``, ``_next_kv_prefetch``); hints only. + """ + from .kernels import m3_attn_lc8 as _attn_lc8_kernel + + core = self.core + backend = core.attn[i] + num_tokens = int(q.shape[0]) + rows = int(step["num_seqs"]) + qlen = int(step["msa_decode_span"].query_len) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + num_kv_heads = int(backend.indexer.config.num_kv_heads) + # The index scorers are compiled and loaded by prebuild_kernels, so no step compiles one and the identity can + # be taken from the first decode step. + identity = step["msa_short_ok"] and step["msa_short_table"] is not None + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + else: + max_score = self._index_scores(i, idx_q, step, num_tokens, qlen) + topk = self._select(i, max_score, step, num_tokens).permute(1, 0, 2) + manager = step["kv_cache_manager"] + k_paged, v_paged = msa_paged_kv(manager, i) + nvfp4 = {} + if manager.is_nvfp4_layer(i): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(backend, core.kv_scales[i]) + scales = manager.get_block_scale_buffers(i, "HND") + nvfp4 = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + # The long-context decode (if it runs) prefetches this layer's o_proj weights into L2; any other route leaves + # the slot unused, and it is cleared after the call either way. + if _attn_lc8_kernel.MERGE_ONCE and SPARSE_DECODE_LC8_PREFETCH_OPROJ: + _attn_lc8_kernel.PREFETCH.tensors = (core.w[f"l{i}_o_w"], core.w[f"l{i}_o_s"]) + if _attn_lc8_kernel.MERGE_ONCE and i in core.n1_ids and num_tokens <= N1_MAX_ROWS: + _attn_lc8_kernel.NEXT_KV.tensors = self._next_kv_prefetch(i, manager) + _record("m3_sparse_decode", num_tokens, qlen, identity) + try: + m3_sparse_decode( + q.view(num_tokens, core.heads, core.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=core.sm_scale(i), + output=output.view(num_tokens, core.heads, core.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4, + ) + finally: + _attn_lc8_kernel.PREFETCH.tensors = None + _attn_lc8_kernel.NEXT_KV.tensors = None + + def _next_kv_prefetch(self, i: int, manager): + """(layer ``i``'s id buffer, the next sparse layer's, that layer's K / V pools and block scales) for the + long-context decode's next-layer prefetch; the last sparse layer's next is the first one (next step). None when + the next layer's cache is not NVFP4.""" + n1_ids = self.core.n1_ids + order = sorted(n1_ids) + nxt = next((j for j in order if j > i), order[0]) + if not manager.is_nvfp4_layer(nxt): + return None + k_next, v_next = msa_paged_kv(manager, nxt) + scales = manager.get_block_scale_buffers(nxt, "HND") + return (n1_ids[i], n1_ids[nxt], k_next, v_next, scales[:, 0], scales[:, 1]) + + +_ENGAGED: set = set() + + +def _log_engaged(path: str, mapping) -> None: + """One info line per fast path and process (every GEN rank logs its own): the run's engagement evidence.""" + if path not in _ENGAGED: + _ENGAGED.add(path) + rank = getattr(mapping, "tp_rank", "?") if mapping is not None else "?" + logger.info(f"[M3 sm_107 tp8] fast path engaged: {path} (tp rank {rank})") + + +class M3SpecStepFold: + """The target's speculative step-input staging (the in-tree ``SpecStepFold`` protocol). + + Between two decode graphs of a one-model Eagle3 step, the spec sampler stores the step's outputs into its per-slot + store and the engine gathers them back into its flat input buffers for the next step. This object does both in + one launch each (``m3_spec_tail.slot_scatter`` / ``step_gather``): integer copies only, the same values the + engine's and the sampler's own copies write. + + The forward of a decode-only step also does both at its tail (``store_in_forward``, one ``m3_fold_tail`` launch + inside the step graph, this target's C++ op): it stores the step's outputs for the rows ``begin_step`` armed and + stages the next step's inputs for those rows in that order. The arming is a device-side row count and slot list + staged before the replay, so a captured graph whose step was not armed writes nothing. The sampler then skips its + store for exactly the armed rows (``sampler_rows_stored``), and the next step skips its gathers when its rows are + those rows (``begin_step``). The sampler copies the store to the host on its D2H side stream and records + ``store_copies_done`` after the copies; the forward's tail waits on it before it rewrites the store. + + * ``store`` -- the spec sampler's per-slot store (``bind_sampler_store``). + * ``step_inputs`` -- the engine's (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) device buffers + (``bind_engine_buffers``). + """ + + # Read by the shell's forward (the tail runs only when the fold stores in the forward). + in_forward = True + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[tuple] = None + # Device arming read by the forward's tail: ctrl[0] = armed rows (0 = write nothing), slots[:rows]. + self._ctrl: Optional[torch.Tensor] = None + self._slots: Optional[torch.Tensor] = None + # The rows (sequence slots) the next forward stores and stages for, as begin_step armed them. + self._armed: Optional[tuple] = None + # The rows the previous forward staged the next step's inputs for (confirmed by sampler_rows_stored). + self._staged_for: Optional[tuple] = None + # Output buffers the tail ran on: the captured graphs' static outputs (by address) and the last eager step's. + self._graph_outputs: set = set() + self._eager_outputs: Optional[torch.Tensor] = None + self._logged = False + # The (external) event the sampler records on its D2H side stream after copying the store to the host, waited + # on by the forward's tail before it rewrites the store (see SpecSamplerBase._copy_outputs_to_host). + self.store_copies_done: Optional[torch.cuda.Event] = None + + def bind_engine_buffers( + self, input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ): + self.step_inputs = (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) + # Recorded once now so a captured wait has a completed record to wait on before the first copy. + self.store_copies_done = torch.cuda.Event(external=True) + self.store_copies_done.record() + logger.info( + "[M3 host] spec store host copies on the D2H side stream, awaited by the forward's tail" + ) + logger.info("[M3 host] spec fold tail: one C++ launch (no Triton) at the forward's tail") + self._ctrl = torch.zeros(1, dtype=torch.int32, device=input_ids.device) + self._slots = torch.zeros(max(max_batch_size, 1), dtype=torch.long, device=input_ids.device) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + if self._ctrl is None: + return False + rows = tuple(seq_slots) if seq_slots is not None else None + staged = rows is not None and rows == self._staged_for + self._staged_for = None + armed = rows if rows is not None and 0 < len(rows) <= self._slots.numel() else None + # Read only by this step's forward, so both uploads may wait for the step-input pack's flush. + if armed is not None: + slots = torch.tensor(armed, dtype=torch.long, pin_memory=prefer_pinned()) + copy_h2d_if_changed(self._slots[: len(armed)], slots, pack=True) + ctrl = torch.tensor( + [len(armed) if armed is not None else 0], dtype=torch.int32, pin_memory=prefer_pinned() + ) + copy_h2d_if_changed(self._ctrl, ctrl, pack=True) + self._armed = armed + return staged + + # The launches the store wait may precede. store_copies_done is an external event: under graph capture its wait + # is a node of its own, so the kernel launched right after it loses its programmatic (PDL) edge and with it the + # overlap of its pre-wait part. Only a launch that waits for its predecessor's completion anyway may follow it: + # the fold tail (a plain launch, no griddepcontrol). Waiting before the last draft_step instead cost +13 / +17 + # us/step. + STORE_WAIT_BEFORE = frozenset({"fold_tail"}) + + def _wait_store_copies(self, before: str) -> None: + """Order the store's rewrite after the sampler's host copies of it; ``before`` names the launch that follows. + Under graph capture a launch outside STORE_WAIT_BEFORE fails the capture.""" + if self.store_copies_done is None: + return + if before not in self.STORE_WAIT_BEFORE and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"M3SpecStepFold: the store wait would precede {before!r}, a launch outside " + f"{sorted(self.STORE_WAIT_BEFORE)}; an external-event wait node drops the next kernel's PDL edge" + ) + torch.cuda.current_stream().wait_event(self.store_copies_done) + + def store_in_forward(self, outputs: dict) -> None: + """At the forward's tail: store the armed rows' outputs and stage the next step's inputs (one launch).""" + if self._ctrl is None or self.store is None or self.step_inputs is None: + return + store = self.store + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + new_tokens = outputs.get("new_tokens") + next_new_tokens = outputs.get("next_new_tokens") + new_tokens_lens = outputs.get("new_tokens_lens") + next_draft_tokens = outputs.get("next_draft_tokens") + if any( + t is None or not isinstance(t, torch.Tensor) + for t in (new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens) + ): + return + rows = new_tokens.shape[0] + # The previous step's host copies of the store run on the sampler's side stream. + self._wait_store_copies(before="fold_tail") + _record("m3_fold_tail", rows) + if not m3_fold_tail.fold_tail( + self._ctrl, + self._slots, + min(rows, self._slots.numel()), + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + next_draft_tokens, + store.next_draft_tokens, + new_tokens_lens, + store.new_tokens_lens, + input_ids, + draft_tokens, + position_id_offsets, + kv_lens_offsets, + next_new_tokens.shape[1], + cpp=True, + ): + return + if torch.cuda.is_current_stream_capturing(): + self._graph_outputs.add(new_tokens.data_ptr()) + else: + self._eager_outputs = new_tokens + + def gather( + self, + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + if self.step_inputs is None: + return False + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + return m3_spec_tail.step_gather( + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len, + tokens_per_gen_step, + draft_token_buffer_width, + input_ids, + num_tokens, + draft_tokens, + num_draft_tokens, + position_id_offsets, + kv_lens_offsets, + first_previous_row, + ) + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list) -> bool: + armed, self._armed = self._armed, None + new_tokens = outputs.get("new_tokens") + ran = isinstance(new_tokens, torch.Tensor) and ( + new_tokens is self._eager_outputs or new_tokens.data_ptr() in self._graph_outputs + ) + self._eager_outputs = None + stored = ( + ran + and armed is not None + and num_skip == 0 + and tuple(seq_slots) == armed + and new_tokens.shape[0] == len(armed) + ) + self._staged_for = armed if stored else None + if stored and not self._logged: + self._logged = True + logger.info("[M3 host] spec step stored and next step staged in the forward tail") + return stored + + def scatter( + self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) -> bool: + store = self.store + if store is None: + return False + return m3_spec_tail.slot_scatter( + seq_slots, + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + new_tokens_lens, + store.new_tokens_lens, + next_draft_tokens, + store.next_draft_tokens, + ) + + +class _DraftForward: + """``draft_model.model`` for the in-tree worker: the drafter layer's forward, called as + ``model(input_ids=, position_ids=, hidden_states=, attn_metadata=, spec_metadata=, all_rank_num_tokens=)`` and + returning (hidden, hidden to carry to the next pass). ``d2t`` None: the draft vocabulary is the target's.""" + + d2t = None + + def __init__(self, draft: "M3Eagle3Draft") -> None: + self._draft = draft + + def __call__(self, **kwargs): + return self._draft.forward_layer(**kwargs) + + +# The drafter's projections by weight-table name, for the low-M linears' engagement and prebuild. +_DRAFT_LINEARS = ("fc", "qkv", "o", "gu", "dn") + + +class M3Eagle3Draft(nn.Module): + """The one-layer BF16 Eagle3 drafter (``speculative_model``): its weights (``weights.DRAFT_WEIGHTS``), its + vocab-parallel BF16 head and logits processor, the drafter layer's TRTLLM attention backend (layer index = the + target's layer count, FP8 KV) and its TP all-reduce. + + What the worker reads: ``model`` (the layer forward), ``apply_eagle3_fc`` (the per-layer fc_norms + fc over the + captured target hidden states), ``lm_head`` and ``logits_processor``; the fused draft loop reads the layer's parts + (``draft_input_weights``, ``attention``, ``mlp``, ``ar_norm``). The drafter embeds with the target's table. + + The BF16 projections (fc, qkv, o_proj, gate_up, down; not the head) run FlashInfer's fixed-tactic low-M kernels + (``ops/gemm/m3_draft_lowm``, the TP8 port) at the rows ``DRAFT_LOWM`` covers, compiled before the warmup + (``prebuild_linears``); every other call is the in-tree cuBLAS GEMM. Rubin TP2 (mclk locked, 4 rotating sets) at + 1 row: qkv 13.18 -> 9.71, o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. The + split-K ('s') tactics run this target's copy of FlashInfer's kernel (``kernels/m3_lowm_splitk``, bit for bit the + same output); with ``DRAFT_AR_PUSH`` the o_proj and down projection of the fused draft loop also push their rows + into the drafter boundary all-reduce's exchange (``arm_push`` / ``take_pushed``); the direct ('d') tactics never + push. + """ + + def __init__( + self, + core: MiniMaxM3Tp8ModelingV2, + draft_config: ModelConfig, + start_layer_idx: int, + capture_layers: int, + ) -> None: + super().__init__() + cfg = draft_config.pretrained_config + _ensure_draft_vocab_size(cfg) + eagle = getattr(cfg, "eagle_config", None) or {} + mapping = draft_config.mapping + self.draft_config = draft_config + # The engine's draft weight mapper binds to these two. + self.model_config = draft_config + self.config = cfg + self.capture_layers = int(capture_layers) + self.layer_idx = int(start_layer_idx) + self.dtype = cfg.torch_dtype + self.hidden = int(cfg.hidden_size) + self.hidden_size_in = int(getattr(cfg, "target_hidden_size", cfg.hidden_size)) + self.rms_eps = float(cfg.rms_norm_eps) + self.return_hidden_post_norm = bool( + eagle.get("return_hidden_post_norm", False) or getattr(cfg, "norm_output", False) + ) + if ( + int(cfg.num_hidden_layers) != 1 + or eagle.get("eh_proj_before_attn", False) + or eagle.get("norm_before_fc", False) + or not getattr(cfg, "fc_norm", False) + or getattr(cfg, "kv_lora_rank", None) + or self.capture_layers <= 1 + or self.hidden_size_in != self.hidden + or cfg.draft_vocab_size != cfg.vocab_size + or getattr(cfg, "attention_bias", False) + or getattr(cfg, "mlp_bias", False) + or getattr(cfg, "bias", False) + or getattr(cfg, "tie_word_embeddings", False) + ): + raise ValueError( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle3 drafter with per-layer " + "fc_norms over " + "several captured layers, the target's hidden size and vocabulary and no biases" + ) + tp = mapping.tp_size + self.heads = cfg.num_attention_heads // tp + self.kv_heads = (cfg.num_key_value_heads + tp - 1) // tp + self.head_dim = int( + getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + ) + self.q_size = self.heads * self.head_dim + self.kv_size = self.kv_heads * self.head_dim + self.w = _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + self.lm_head = LMHead( + cfg.draft_vocab_size, + cfg.hidden_size, + dtype=cfg.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=draft_config.lm_head_gather_output, + reduce_output=False, + use_custom_cublas_mm=False, + quant_config=None, + ) + self.logits_processor = LogitsProcessor() + self.allreduce = AllReduce( + mapping=mapping, strategy=draft_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + quant_config = draft_config.get_quant_config() + self.attn = create_attention( + draft_config.attn_backend, + self.layer_idx, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, rope=RopeParams.from_config(cfg) + ), + quant_config=quant_config, + skip_create_weights_in_init=draft_config.skip_create_weights_in_init, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): + raise ValueError( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + # The drafter embeds with the target's table, read off the core on every call (the engine materializes the + # core's weights after construction). + self._core = [core] + self.model = _DraftForward(self) + # The low-M linears: (N, K, M, PDL) of every call prebuild_linears compiled, and (N, K, M) of every call that + # logged its engagement line. + self._lowm_built: set = set() + self._lowm_engaged: set = set() + # DRAFT_AR_PUSH: the weight whose next low-M call also pushes its rows into a boundary all-reduce (data + # pointer, the bucket's multicast words, its flags, rank, world, max tokens), set by arm_push; and the data + # pointer of the output the last armed call pushed, 0 if none (read once by take_pushed). Host state set and + # read in one forward's issue order. + self._push_arm: Optional[tuple] = None + self._pushed = 0 + + @property + def embed_table(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def load_weights_from_target_model(self, target_model: nn.Module) -> None: + """Nothing to alias: the drafter reads the target's embedding table through the core and owns its head.""" + + # ------------------------------------------------------------------------------------------- the projections + + def linear(self, x: torch.Tensor, key: str) -> torch.Tensor: + """``x @ w[key]^T``: the low-M kernel where ``DRAFT_LOWM`` covers the call, else the in-tree cuBLAS GEMM.""" + w = self.w[key] + out = self._lowm_linear(w, x) + return out if out is not None else m3_bf16_linear(x, w) + + def _lowm_linear(self, weight: torch.Tensor, x: torch.Tensor) -> Optional[torch.Tensor]: + """x [M, K] @ weight^T through the low-M kernel when ``DRAFT_LOWM`` covers it; None for the in-tree GEMM. + A projection none of whose row counts is covered never reaches the low-M checks.""" + n, k = weight.shape + if x.dim() != 2 or not any(nk == (n, k) for nk, _ in DRAFT_LOWM.keys): + return None + m = int(x.shape[0]) + if not DRAFT_LOWM.covers(((n, k), m)): + uncovered(DRAFT_LOWM.path, f"the drafter's {n}x{k} linear past its tactic table's rows") + return None + if not x.is_contiguous(): + # The first pass's fc may read a column window of the captured hidden states; the in-tree GEMM takes + # strides. + uncovered(DRAFT_LOWM.path, "a strided activation (the in-tree GEMM reads it in place)") + return None + use_pdl = get_env_enable_pdl() + if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + missed(DRAFT_LOWM.path, m, "the activation or weight is not BF16") + return None + if (n, k, m, use_pdl) not in self._lowm_built: + missed( + DRAFT_LOWM.path, m, f"the {n}x{k} kernel at {m} rows was not compiled before warmup" + ) + return None + kind, tactic = tactic_for(n, k, m) + if (n, k, m) not in self._lowm_engaged: + self._lowm_engaged.add((n, k, m)) + logger.info( + f"[M3 sm_107 tp8] fast path engaged: m3_draft_lowm {n}x{k} at {m} rows (device {x.device.index})" + ) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + ar = None + if ( + self._push_arm is not None + and self._push_arm[0] == weight.data_ptr() + and kind == "s" + and DRAFT_LOWM_OWN_SPLITK + ): + ar = self._push_arm[1:] + self._pushed = out.data_ptr() + _record("m3_draft_lowm", n, k, m, kind, ar is not None) + return m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, DRAFT_LOWM_OWN_SPLITK, ar) + + def prebuild_linears(self) -> int: + """Compile (and run once) every covered row count of every drafter projection, outside capture, on the + current stream, before the executor's warmup. Returns the number of calls compiled.""" + from .kernels.m3_lowm_splitk import compile_push + + use_pdl = get_env_enable_pdl() + built = 0 + shapes = {tuple(self.w[key].shape): self.w[key] for key in _DRAFT_LINEARS} + for (n, k), weight in sorted(shapes.items()): + if weight.dtype != torch.bfloat16: + continue + for _, m in sorted(key for key in DRAFT_LOWM.keys if key[0] == (n, k)): + if (n, k, m, use_pdl) in self._lowm_built: + continue + kind, tactic = tactic_for(n, k, m) + x = torch.zeros(m, k, dtype=torch.bfloat16, device=weight.device) + out = torch.empty(m, n, dtype=torch.bfloat16, device=weight.device) + m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, DRAFT_LOWM_OWN_SPLITK) + if DRAFT_AR_PUSH and DRAFT_LOWM_OWN_SPLITK and kind == "s": + # The push variant a boundary all-reduce may arm (o_proj / down): compiled only, never run here + # (its pushes need the all-reduce's exchange and a matching call). + compile_push(x, weight.t(), out, use_pdl, tactic) + self._lowm_built.add((n, k, m, use_pdl)) + built += 1 + return built + + def arm_push(self, key: str, workspace, rows: int) -> None: + """Arm ``w[key]``'s next covered low-M call to push into ``workspace``'s bucket for ``rows``.""" + bucket = m3_mc_allreduce._bucket(workspace, rows) + flags = None if bucket is None else bucket.flags_for(8) + self._pushed = 0 + self._push_arm = None + if bucket is not None and flags is not None and bucket.max_tokens == 8: + self._push_arm = ( + self.w[key].data_ptr(), + bucket.mc.data_ptr(), + flags.data_ptr(), + bucket.rank, + bucket.world_size, + bucket.max_tokens, + ) + + def take_pushed(self, x: torch.Tensor) -> bool: + """Whether the armed call pushed its rows; disarms. ``x`` is what the boundary all-reduce is about to reduce: + it must be the pushing call's own output (the prepushed all-reduce reads only the pushed words, never ``x``), + so an operation between that call and the all-reduce raises instead of silently reducing the wrong rows.""" + pushed, self._pushed, self._push_arm = self._pushed, 0, None + if pushed and x.data_ptr() != pushed: + raise RuntimeError( + "the drafter's pushed rows are not the boundary all-reduce's input (an operation ran in between)" + ) + return bool(pushed) + + # ------------------------------------------------------------------------------------------------- the layer + + def fc(self, hs: torch.Tensor) -> torch.Tensor: + return self.linear(hs, "fc") + + def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The captured target hidden states [T, layers x H] to the drafter's input width: each layer's fc_norm, + then the fc (the in-tree ``Eagle3ForCausalLM.apply_eagle3_fc``).""" + hidden_states = hidden_states.to(self.dtype) + if hidden_states.shape[-1] != self.hidden: + chunks = hidden_states.chunk(self.capture_layers, dim=-1) + hidden_states = torch.cat( + [ + flashinfer_rmsnorm(c, self.w[f"fc_norm{i}"], self.rms_eps) + for i, c in enumerate(chunks) + ], + dim=-1, + ) + hidden_states = self.fc(hidden_states) + return hidden_states + + def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: + """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is + False, where the caller fuses the reduction). The backend's generation phase runs this target's full-context + decode with the fused RoPE + FP8 K/V append when the shell installed it (``kernels/m3_drafter_attn``).""" + qkv = self.linear(x, "qkv") + _record("m3_draft_attention", int(md.num_tokens)) + o = m3_draft_attention(self.attn, qkv, md) + out = self.linear(o, "o") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool, cpp_silu: bool = False) -> torch.Tensor: + """gate_up, the plain SwiGLU, down (all-reduced unless ``reduce`` is False). ``cpp_silu`` (the fused draft + loop): the SwiGLU on this target's C++ op (``ops/spec/m3_draft_silu``, minimaxM3Sm107Tp8/siluAndMul.cu), bit + for bit the in-tree Triton silu_and_mul in one launch either way (no Triton on the GEN decode path).""" + gu = self.linear(h, "gu") + act = None + if cpp_silu: + act = m3_draft_silu(gu) + if act is None: + uncovered( + DRAFT_SILU_PATH, "the drafter SwiGLU input layout (the Triton op runs it)" + ) + else: + taken(DRAFT_SILU_PATH, int(gu.shape[0]), "drafter MLP SwiGLU, one C++ launch") + if act is None: + act = m3_silu_and_mul(gu, None, None, None) + out = self.linear(act, "dn") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def ar_norm(self, x: torch.Tensor, residual: torch.Tensor, norm_key: str): + """One draft-layer boundary in the in-tree structure: the projection's plain all-reduce (BF16 out), then the + fused residual add + RMSNorm (FP32 add, norm of the FP32 sum). Returns (normed, residual sum).""" + return flashinfer_fused_add_rmsnorm( + m3_allreduce(x, self.allreduce, None), residual, self.w[norm_key], self.rms_eps + ) + + def forward_layer( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + all_rank_num_tokens: Optional[List[int]] = None, + ): + """The drafter layer (the in-tree ``Eagle3DraftModel.forward`` for one layer): the embedded tokens and the + fc'd hidden states normed and concatenated, attention, the post-attention norm with the residual, the MLP and + the final norm. Returns (normed output, the hidden state the next pass carries).""" + previous = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + embeds = m3_embedding(input_ids, self.embed_table).to(self.dtype) + residual = hidden_states + h = flashinfer_rmsnorm(hidden_states, self.w["hidden_norm"], self.rms_eps) + e = flashinfer_rmsnorm(embeds, self.w["in_norm"], self.rms_eps) + h = torch.cat([e, h], dim=-1) + h = self.attention(h, attn_metadata, reduce=True) + h, residual = flashinfer_fused_add_rmsnorm( + h, residual, self.w["post_norm"], self.rms_eps + ) + h = self.mlp(h, reduce=True) + spec_metadata.maybe_capture_hidden_states(self.layer_idx, h, residual) + h, to_save = flashinfer_fused_add_rmsnorm( + h, residual, self.w["final_norm"], self.rms_eps + ) + return (h, h) if self.return_hidden_post_norm else (h, to_save) + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous + + +def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, capture_norm) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "tp_size", 1) > 1 + ) + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify and fused draft loop, at the in-tree numerics. + + * Verify: on decode-only batches whose verify rows the fused kernels cover (``tail_eligible``), the target's BF16 + lm_head computes this rank's vocabulary shard of the logits (the in-tree head's own local GEMM) and one kernel + replaces the logits all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. + * Draft: the drafter's first pass may take the per-layer fc_norms in one launch (or folded into the hidden-state + capture) and run its fc beside the target head GEMM (forked right after the last captured target layer at most + ``FC_EARLY_MAX_ROWS`` verify rows; after a sampled verify's head at most ``FC_AFTER_HEAD_MAX_ROWS``). On a + decode-only step ``DRAFT_STEP`` covers, the draft loop is fused: each pass's draft input in one launch, the + drafter layer's attention and MLP with their all-reduces fused with the residual add and the next RMSNorm (the + multicast all-reduce at ``DRAFT_MC_AR``'s rows, fed by the projection's push), and one kernel for the BF16 head, + the cross-rank argmax and the pass's bookkeeping. + + Every other batch takes the in-tree paths. Methods whose docstring says "Mirrors" copy the named in-tree method + with the fused branches inserted. + """ + + def __init__(self, spec_config, mapping, use_separate_draft_kv_cache: bool = False) -> None: + super().__init__( + spec_config, mapping=mapping, use_separate_draft_kv_cache=use_separate_draft_kv_cache + ) + # Per-step state, reset at the top of every _forward_impl: the first draft pass's gather rows when this step + # takes the fused first drafter inputs, and the fc output a fork computed beside the target head. + self._m3_first_gather_ids: Optional[torch.Tensor] = None + self._m3_fc_hidden: Optional[torch.Tensor] = None + self._m3_target_lm_head: Optional[nn.Module] = None + # (fc_norm0 weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple] = None + # Per step: the captured layers written with the capture norm (set by capture_hidden_states). + self._m3_capture_normed: set = set() + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + # Armed per step by arm_capture_norm: the drafter whose first fc is forked at the last captured target layer, + # and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + # DI_FC_FLAG: int64 [3] device counters (forks bumped, forks waited, wait overruns), allocated by ``prebuild``; + # whether the latest fork bumped them; the fork whose stream join is deferred past the first draft input. + self._m3_fc_counts: Optional[torch.Tensor] = None + self._m3_fc_flagged = False + self._m3_fc_pending: Optional[torch.cuda.Stream] = None + # The 12-bit coded copy of the target head shard the fused verify streams, built before the warmup. + self._m3_force_vec_warm = False + # The streamed copy of the drafter's BF16 head shard the fused draft loop reads, built before the warmup, the + # head weight version it was built from, and the replaced copies a captured graph may still read. + self._m3_draft_head = None + self._m3_draft_head_version: Optional[int] = None + self._m3_draft_head_graveyard: list = [] + # Per step: the next step's input tokens [nreq, D + 1], column 0 written by the fused verify (or the seed) and + # columns 1.. by the fused draft passes; set as _m3_next_new_tokens once the fused draft loop filled it. + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + # Per step: the first draft inputs are the fused ones (accepted tokens as they lie, fc output of every row). + self._m3_fused_first_inputs = False + + def capture_target(self, spec_metadata, layer_id: int, num_tokens: int) -> Optional[tuple]: + """Where this step's capture norm writes layer ``layer_id``'s capture, for a boundary all-reduce that writes it + itself (CAPTURE_NORM_IN_BOUNDARY): (the capture buffer's [num_tokens, H] column block, the layer's fc_norm + weight, eps), or None (no capture norm armed, not a captured layer, or more rows than the buffer). The + boundary writes every row it reduces; rows past num_capture_tokens are unread.""" + capture_norm = self._m3_capture_norm + if capture_norm is None or spec_metadata.hidden_states is None: + return None + if not m3_mc_allreduce.capture_supported(): + return None + if num_tokens > spec_metadata.hidden_states.shape[0]: + return None + hidden_size = spec_metadata.hidden_size + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + out = spec_metadata.hidden_states[ + :num_tokens, i * hidden_size : (i + 1) * hidden_size + ] + return out, capture_norm[0][i], capture_norm[1] + return None + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None, written: bool = False + ) -> None: + """The Eagle3 capture with this step's capture norm (``written``: the layer's boundary all-reduce already wrote + it, see ``capture_target``); on a step armed for the early fc fork (``arm_capture_norm``), the last captured + layer then forks the drafter's first fc onto the aux stream, where it overlaps the remaining target layers + and the head (joined by ``_m3_target_accept``).""" + if not written: + _eagle3_capture(spec_metadata, layer_id, hidden_states, residual, self._m3_capture_norm) + if self._m3_capture_norm is not None: + # Written with this step's capture norm: by the boundary all-reduce (written) or just above. + self._m3_capture_normed.add(layer_id) + drafter = self._m3_fc_early_drafter + layers = spec_metadata.layers_to_capture + if drafter is not None and layers and layer_id == layers[-1]: + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = self._m3_fc_fork( + spec_metadata, spec_metadata.num_capture_tokens, drafter + ) + + def prebuild(self, lm_head: nn.Module, draft: Optional[M3Eagle3Draft]) -> None: + """From the loaded weights, before the executor's warmup, on every rank: the fused draft loop's head copy and + the drafter's low-M linears.""" + start = time.perf_counter() + if self.prebuild_draft_head(draft): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + if DI_FC_FLAG and draft is not None and self._m3_fc_counts is None: + # The fc fork counters (DI_FC_FLAG), zeroed here, before warmup, in order before any fork. + self._m3_fc_counts = torch.zeros( + 3, dtype=torch.int64, device=draft.w["hidden_norm"].device + ) + start = time.perf_counter() + built = draft.prebuild_linears() if draft is not None else 0 + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] drafter low-M GEMMs built before warmup: %d calls, %.1f s" + % (built, time.perf_counter() - start) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Mirrors Eagle3OneModelWorker._forward_impl, with the fused verify branch.""" + runtime_draft_len = spec_metadata.runtime_draft_len + # skip the draft forward if the runtime draft length is 0 + if runtime_draft_len == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._m3_first_gather_ids = None + self._m3_next_new = None + self._m3_next_new_tokens = None + self._m3_fused_first_inputs = False + if logits is None: + # tail_eligible() skipped the gathered target logits: verify from this rank's logits shard. + accepted_tokens, num_accepted_tokens, self._m3_first_gather_ids = ( + self._m3_target_accept( + hidden_states, + spec_metadata, + batch_size, + runtime_draft_len, + draft_model, + kv_lens=getattr(attn_metadata, "kv_lens_cuda", None), + ) + ) + else: + self._execute_guided_decoder_if_present(logits) + + # Sample and accept tokens. ``input_ids`` is required by the relaxed-acceptance path (scans for + # thinking-phase tokens); ignored otherwise. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, attn_metadata, spec_metadata + ) + if self.draft_only_eligible(attn_metadata, spec_metadata, draft_model): + self._m3_seed_first_pass(num_accepted_tokens, spec_metadata, batch_size) + + # Mamba hybrid models need state updates after token acceptance (a no-op on this target's cache manager). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None) + if self.sa_enhancer is not None and sa_manager is not None: + self.sa_enhancer.extend_and_prepare( + sa_manager=sa_manager, + request_ids=spec_metadata.request_ids, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_gens=num_gens, + num_contexts=num_contexts, + max_draft_len=runtime_draft_len, + ) + + if self._uses_external_shared_target_kv: + next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + ) + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here so the post-loop restore (below) + # can put attn_metadata back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + self._m3_join_pending() + # prepare next new tokens to support overlap scheduler + if self._m3_next_new_tokens is not None: + # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). + next_new_tokens = self._m3_next_new_tokens + else: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + # Commit model-owned target state only after draft preparation has succeeded; SpecWorkerBase aborts the + # still-pending snapshots on any exception above this point. + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices[:batch_size], + num_contexts, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """The fused first draft inputs on a step that verified with the fused kernels (or was seeded for them), + else the in-tree ones.""" + fused = self._m3_first_drafter_inputs( + position_ids, + accepted_tokens, + input_ids.shape[0], + attn_metadata.num_contexts, + spec_metadata, + draft_model, + ) + self._m3_fused_first_inputs = fused is not None + if fused is not None: + fused.update(attn_metadata=attn_metadata, spec_metadata=spec_metadata) + return fused + return super().prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + def _m3_structural_ok(self, attn_metadata, spec_metadata) -> bool: + """The batch conditions shared by the fused verify and the fused first drafter inputs: a decode-only plain + Eagle3 step without attention DP, guided decoding, relaxed acceptance or a suffix-automaton enhancer, of at + most DECODE_MAX_TOKENS verify rows. Host state only, so the TP ranks agree.""" + if self.is_mtp_eagle or self.use_dynamic_tree: + return False + if self._uses_external_shared_target_kv or self.sa_enhancer is not None: + return False + if ( + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp + ): + return False + if getattr(self.spec_config, "use_relaxed_acceptance_for_thinking", False): + return False + if attn_metadata.num_contexts != 0 or attn_metadata.num_seqs == 0: + return False + draft_len = spec_metadata.runtime_draft_len + if draft_len == 0 or attn_metadata.num_seqs * (draft_len + 1) > DECODE_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Whether this step verifies with the fused kernels instead of the gathered target logits. + + Decode-only Eagle3 batches (``_m3_structural_ok``) over an unpadded vocab-sharded bias-free target head, all + greedy at a request count ``SHARD_ACCEPT`` covers, or sampled under the temperature + top-p contract (see + _m3_sampled_verify_ok) at a request count ``SAMPLED_ACCEPT`` covers, with nothing else reading the target + logits. Every input is host state shared by all TP ranks, so the ranks agree. + """ + if not self._m3_structural_ok(attn_metadata, spec_metadata): + return False + if getattr(spec_metadata, "enable_penalty", False): + return False + if not _sharded_head(lm_head) or getattr(lm_head, "bias", None) is not None: + return False + draft_len = spec_metadata.runtime_draft_len + if spec_metadata.is_all_greedy_sample: + key = (draft_len, attn_metadata.num_seqs) + if not SHARD_ACCEPT.covers(key): + uncovered( + SHARD_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + else: + if not self._m3_sampled_verify_ok(spec_metadata): + return False + key = (draft_len, attn_metadata.num_seqs) + if not SAMPLED_ACCEPT.covers(key): + uncovered( + SAMPLED_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + self._m3_target_lm_head = lm_head + _log_engaged( + SHARD_ACCEPT.path if spec_metadata.is_all_greedy_sample else SAMPLED_ACCEPT.path, + self.mapping, + ) + return True + + def draft_only_eligible(self, attn_metadata, spec_metadata, draft_model) -> bool: + """Whether a step verified by the in-tree path can still take the fused first drafter inputs: up to + FIRST_DRAFTER_INPUTS_MAX_REQUESTS requests, the structural conditions of tail_eligible, without the + target-sampling ones.""" + if attn_metadata.num_seqs > FIRST_DRAFTER_INPUTS_MAX_REQUESTS: + uncovered( + "m3_first_drafter_inputs", f"in-tree verify at {attn_metadata.num_seqs} requests" + ) + return False + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None + + def _m3_seed_first_pass(self, num_accepted_tokens, spec_metadata, batch_size: int) -> None: + """The first draft pass's gather rows (the last accepted row of each request) from the in-tree acceptance, as + _m3_target_accept leaves them.""" + k1 = spec_metadata.runtime_draft_len + 1 + n = num_accepted_tokens[:batch_size].to(torch.int64) + rows = spec_metadata.batch_indices_cuda[:batch_size].to(torch.int64) + self._m3_first_gather_ids = rows * k1 + n - 1 + + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], + ): + """The in-tree first draft inputs of a decode-only step, with the per-layer fc_norm RMSNorms of + ``apply_eagle3_fc`` in one kernel (unless the capture already applied them) before the fc; the input ids are + the accepted tokens as they lie. None to use the in-tree path.""" + if ( + self._m3_first_gather_ids is None + or self.is_mtp_eagle + or num_contexts != 0 + or draft_model is None + ): + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + ): + return None + if not self._m3_first_inputs_covered(num_tokens, spec_metadata): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not draft_model.w["fc_norm0"]: + if torch.cuda.is_current_stream_capturing(): + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), + ) + self._m3_fc_norm_w = cache + hidden = self._m3_fc_hidden + self._m3_fc_hidden = None + if hidden is None or hidden.shape[0] != num_tokens: + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + _record("m3_chunk_norm", num_tokens) + hs = m3_spec_tail.chunk_norm(hs, cache[1], draft_model.rms_eps) + _log_engaged(CHUNK_NORM.path, self.mapping) + else: + hs = self._m3_renorm_raw_captures(hs, spec_metadata, draft_model) + hidden = draft_model.fc(hs) + input_ids = accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].reshape(-1) + _log_engaged(FUSED_FIRST_INPUTS.path, self.mapping) + return {"input_ids": input_ids, "position_ids": position_ids, "hidden_states": hidden} + + def _m3_first_inputs_covered(self, num_tokens: int, spec_metadata) -> bool: + """Whether the fused first drafter inputs of ``num_tokens`` verify rows are covered: their request count + (``FUSED_FIRST_INPUTS``), and the chunk norm's rows (``CHUNK_NORM``) unless the captures already applied the + fc_norm.""" + k1 = spec_metadata.runtime_draft_len + 1 + if num_tokens % k1 or not FUSED_FIRST_INPUTS.covers(num_tokens // k1): + uncovered( + FUSED_FIRST_INPUTS.path, + f"{num_tokens} verify rows at draft length {k1 - 1} (the in-tree first drafter inputs run)", + ) + return False + if self._m3_capture_norm is None and not CHUNK_NORM.covers(num_tokens): + uncovered( + CHUNK_NORM.path, f"{num_tokens} verify rows (the in-tree first drafter inputs run)" + ) + return False + return True + + def _m3_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + after: Optional[torch.cuda.Event] = None, + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (after the event + ``after`` when given, else after everything queued on the current stream); the result is picked up by + _m3_first_drafter_inputs. Returns the stream to join, or None when the fc path does not apply (the first + drafter inputs then compute it).""" + self._m3_fc_hidden = None + if draft_model is None: + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + or self._m3_fc_norm_w is None + or not self._m3_first_inputs_covered(num_tokens, spec_metadata) + ): + return None + aux = self._m3_aux_stream + if aux is None: + aux = self._m3_aux_stream = torch.cuda.Stream() + if after is not None: + aux.wait_event(after) + else: + aux.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(aux): + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + _record("m3_chunk_norm", num_tokens) + hs = m3_spec_tail.chunk_norm( + hs, self._m3_fc_norm_w[1], draft_model.rms_eps, use_pdl=False + ) + _log_engaged(CHUNK_NORM.path, self.mapping) + else: + hs = self._m3_renorm_raw_captures(hs, spec_metadata, draft_model) + self._m3_fc_hidden = draft_model.fc(hs) + self._m3_fc_flagged = self._m3_fc_bump(hs.device) + return aux + + def _m3_renorm_raw_captures( + self, hs: torch.Tensor, spec_metadata, draft_model: M3Eagle3Draft + ) -> torch.Tensor: + """The captured rows ``hs`` [rows, layers * H] for a reader that skips the fc_norm (the capture norm is armed): + unchanged when every captured layer was written with this step's norm (capture_hidden_states records them), + else -- not expected on any path -- a copy with each layer captured raw normalized by its own fc_norm here, + logged as a miss, so a raw row never reaches the drafter's fc.""" + missing = [ + i + for i, layer_id in enumerate(spec_metadata.layers_to_capture) + if layer_id not in self._m3_capture_normed + ] + if not missing: + return hs + missed( + CAPTURE_NORM.path, + int(hs.shape[0]), + f"captured layer(s) {[spec_metadata.layers_to_capture[i] for i in missing]} were written without the " + "capture norm; their fc_norm is applied before the drafter's fc", + ) + hs = hs.clone() + hidden = spec_metadata.hidden_size + for i in missing: + block = hs[:, i * hidden : (i + 1) * hidden] + block.copy_( + flashinfer_rmsnorm(block, draft_model.w[f"fc_norm{i}"], draft_model.rms_eps) + ) + return hs + + def _m3_fc_bump(self, device: torch.device) -> bool: + """DI_FC_FLAG: count the fc just queued on the aux stream (the wait kernel's release). The counters are zeroed + by ``prebuild``; without them the fork keeps the stream join.""" + if not DI_FC_FLAG or self._m3_fc_counts is None or self._m3_fc_counts.device != device: + return False + self._m3_fc_counts[:1].add_(1) + return True + + def _m3_fc_join(self, fork: torch.cuda.Stream) -> None: + """Join the fc fork onto the current stream and keep its output's memory alive for that stream. + + _m3_fc_hidden is allocated in fork (aux) stream order and read on the current stream. Without the record, + dropping it returns its block to the aux stream's pool, where later aux work can take it while the current + stream's reads are still queued (eager warm-up / mixed steps). Nothing is recorded under graph capture, as in + maybe_execute_in_parallel: a captured graph relies on the next fork's aux.wait_stream(main) to order the + reuse.""" + main = torch.cuda.current_stream() + if self._m3_fc_flagged: + # DI_FC_FLAG: the first draft input waits on the fork's counter; the stream join follows that input. + self._m3_fc_flagged = False + self._m3_join_pending() + m3_fc_wait.fc_wait(self._m3_fc_counts) + self._m3_fc_pending = fork + _log_engaged(DI_FC_FLAG_PATH, self.mapping) + else: + main.wait_stream(fork) + hidden = self._m3_fc_hidden + if ( + hidden is not None + and hidden.is_cuda + and not (torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()) + ): + hidden.record_stream(main) + + def _m3_join_pending(self) -> None: + """Join the fc fork whose join DI_FC_FLAG deferred (its output's readiness already went through the wait + kernel): after the first draft input, at the end of the step, and before any newer fork is waited for.""" + fork = self._m3_fc_pending + if fork is not None: + self._m3_fc_pending = None + torch.cuda.current_stream().wait_stream(fork) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model: Optional[M3Eagle3Draft] + ) -> None: + """Before every target forward: whether this step's Eagle3 captures apply the drafter's fc_norm themselves. + + On a fused-verify step of ``CAPTURE_NORM``'s rows whose first drafter inputs are the fused ones and whose + fc_norm weights are already stacked (by an earlier fused first pass), the capture writes + rmsnorm(h + r) * fc_norm[i] (one kernel per captured layer instead of add + copy, or the layer's boundary + all-reduce itself, ``capture_target``) and the fused first drafter inputs read it without their own + chunk_norm. A fused-verify step of at most ``FC_EARLY_MAX_ROWS`` verify rows is also armed to fork the + drafter's first fc at the last captured layer (``capture_hidden_states``). + """ + if self._m3_fc_early_stream is not None: + # A fork the last step made but did not reach the verify: its wait keeps the counters paired. + self._m3_fc_join(self._m3_fc_early_stream) + self._m3_join_pending() + self._m3_capture_norm = None + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = None + # The captured layers this step wrote with its capture norm (capture_hidden_states). + self._m3_capture_normed = set() + # A fork result no first draft pass consumed must not reach a later step's (same-size) first inputs. + self._m3_fc_hidden = None + if draft_model is None: + return + if not self.tail_eligible(attn_metadata, spec_metadata, lm_head): + return + rows = attn_metadata.num_seqs * (spec_metadata.runtime_draft_len + 1) + # The fork reads every captured column at the last captured layer, so a final-state capture, which is written + # after the model returns, keeps the fc at its stock place. + if ( + rows <= FC_EARLY_MAX_ROWS + and FINAL_HIDDEN_LAYER_ID not in spec_metadata.layers_to_capture + ): + self._m3_fc_early_drafter = draft_model + if not CAPTURE_NORM.covers(rows): + uncovered(CAPTURE_NORM.path, f"{rows} verify rows (the in-tree capture runs)") + return + # The normed captures are read only by the fused first drafter inputs; the in-tree ones norm again. + if not FUSED_FIRST_INPUTS.covers(attn_metadata.num_seqs): + return + fc_norm_w = self._m3_fc_norm_w + if ( + fc_norm_w is not None + and fc_norm_w[0] is draft_model.w["fc_norm0"] + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.rms_eps) + _log_engaged(CAPTURE_NORM.path, self.mapping) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata) -> bool: + """Whether sampled verify rows can use the fused sampled op: only under advanced_sampling_mode=no_topk, the + deploy-time guarantee that no row uses top-k or min-p (that op applies neither).""" + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + return mode is not None and mode.skips_top_k and not mode.skips_top_p + + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + kv_lens: Optional[torch.Tensor] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path. ``kv_lens``: the + requests' KV lengths at this verify (the per-request forced-acceptance key, FORCE_AL_PER_REQUEST).""" + force_total = 0 + force_vec = None + if self.force_num_accepted_tokens != 0.0: + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # The draw pool is built by an H2D copy and the draw kernels JIT-compile per block width, so both + # happen here, before the first captured step (a scratch counter keeps the real draw sequence). + self._ensure_force_accept_rng_state(hidden_states.device) + if not self._m3_force_vec_warm: + scratch = torch.zeros(1, dtype=torch.int64, device=hidden_states.device) + dummy = torch.empty(64, dtype=torch.int32, device=hidden_states.device) + for block in (16, 32, 64): + m3_force_vec( + self._force_accept_rng_pool, + scratch, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + m3_force_vec_req( + self._force_accept_rng_pool, + dummy, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + self._m3_force_vec_warm = True + if not is_warmup: + force_total = min(int(self.force_num_accepted_tokens) + 1, draft_len + 1) + force_vec = self._m3_force_vec( + batch_size, force_total, draft_len + 1, hidden_states.device, kv_lens + ) + draft_tokens = spec_metadata.draft_tokens.view(batch_size, draft_len) + if draft_tokens.dtype != torch.int32: + draft_tokens = draft_tokens.to(torch.int32) + # Column 0 receives the last accepted token, the next step's first input (the fused draft loop adds the rest). + self._m3_next_new = torch.empty( + (batch_size, draft_len + 1), dtype=torch.int32, device=hidden_states.device + ) + rows = hidden_states[: batch_size * (draft_len + 1)] + lm_head = self._m3_target_lm_head + batch_idx = spec_metadata.batch_indices_cuda[:batch_size] + # The Eagle3 fc on the aux stream: the fork made at the last captured layer, else one beside the head. A + # sampled verify's head (the in-tree GEMM up to FC_AFTER_HEAD_MAX_ROWS verify rows) holds + # every SM while it streams: its fc is forked after it instead (overlapping the sampled acceptance), where its + # many split-K CTAs do not delay the head's start. Rubin TP2 c1: the head started 17.4 us after its input was + # ready with the fc beside it. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + head_first = ( + fork is None + and rows.shape[0] <= FC_AFTER_HEAD_MAX_ROWS + and not spec_metadata.is_all_greedy_sample + ) + if fork is None and not head_first: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the head's own local GEMM, whose all-gather the fused + # kernels replace (they exchange per-row records across ranks, not logits). + logits = lm_head.apply_linear(rows, None) + head_done = None + if head_first: + head_done = torch.cuda.Event() + head_done.record() + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept_kernel.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = SHARD_ACCEPT.path + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept_kernel.sampled_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + spec_metadata.temperatures[:n_rows], + spec_metadata.top_ks[:n_rows], + spec_metadata.top_ps[:n_rows], + None, + seeds, + offsets, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = "spec_tail" + if head_done is not None: + # Enqueued after the sampled acceptance (which starts first) but waits on the head only. + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model, after=head_done) + if fork is not None: + self._m3_fc_join(fork) + if result is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, None, "the fused verify's exchange workspace") + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + return result + + def _m3_force_vec( + self, + batch_size: int, + base_total: int, + max_total: int, + device, + kv_lens: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + With ``kv_lens`` (FORCE_AL_PER_REQUEST): each request draws from the fixed-seed pool at its own key, its KV + length at this verify (its prompt plus what it has accepted so far), so its draws do not depend on the + server's history or on the other requests; every TP rank holds the same lengths, so the ranks agree. Without + lengths: exactly as _apply_force_accepted_tokens draws for a decode-only batch (one shared counter bump per + step, the (counter, slot) hash), the generic path's sequence. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: + return None + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + block = max(16, 1 << (batch_size - 1).bit_length()) + if FORCE_AL_PER_REQUEST and kv_lens is not None: + keys = kv_lens[:batch_size] + if keys.dtype != torch.int32: + keys = keys.to(torch.int32) + return m3_force_vec_req( + self._force_accept_rng_pool, + keys, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + # One launch: bump the shared counter and draw, with _apply_force_accepted_tokens' exact arithmetic. + return m3_force_vec( + self._force_accept_rng_pool, + self._force_accept_rng_counter, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[M3Eagle3Draft]) -> bool: + """Build the copy of the drafter's BF16 head shard that the fused draft loop streams (rebuilt only when the + head weight tensor or its version changed; a replaced copy is kept alive), on every rank. False when the + drafter is not one the loop serves.""" + old = self._m3_draft_head + if draft_model is None or not self._m3_drafter_supported(draft_model): + self._m3_draft_head = None + return False + weight = draft_model.lm_head.weight + # Idempotent: the same weight tensor at the same version keeps the built copy, so a repeated + # cache_derived_state never re-places it under a captured CUDA graph. + if ( + old is not None + and old.weight is weight + and self._m3_draft_head_version == weight._version + ): + return True + if old is not None: + # A rebuilt copy (new weight values): keep the old one alive, since a graph captured earlier may still + # read its (placed) memory until the graphs are recaptured. + self._m3_draft_head_graveyard.append(old) + self._m3_draft_head = m3_draft_step_kernel.prepare_head( + draft_model.lm_head, self.mapping.tp_rank + ) + self._m3_draft_head_version = weight._version + return True + + def _m3_drafter_supported(self, draft_model: M3Eagle3Draft) -> bool: + """The drafter the fused draft loop runs (``M3Eagle3Draft`` validates its shape at construction) over this + worker's tensor-parallel group with a bias-free vocab-sharded BF16 head without padding.""" + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if self._d2t is not None: + return False + lm_head = draft_model.lm_head + return _sharded_head(lm_head) and m3_draft_step_kernel.head_supported(lm_head) + + def _m3_draft_ar_norm( + self, + draft_model: M3Eagle3Draft, + x: torch.Tensor, + residual: torch.Tensor, + norm_key: str, + carveout: bool = True, + pushed: bool = False, + ): + """One drafter boundary of the fused draft loop: the multicast all-reduce + residual + RMSNorm with the + drafter's stock rounding (``DRAFT_MC_AR``: opts bit 16, normalizing the unrounded FP32 sum as FlashInfer + fused_add_rmsnorm does; Rubin TP2 10.0..11.0 -> 4.6..5.3 us per boundary), else the in-tree structure + (``M3Eagle3Draft.ar_norm``). ``carveout`` False drops the reduce's max-smem carveout (``DRAFT_MC_AR_CARVEOUT``). + ``pushed``: the boundary's projection already pushed its rows into the exchange (prepushed 2). The choice + depends on shapes and a collectively created workspace only, so the TP ranks agree. Returns (normed, residual + sum).""" + rows = x.shape[0] + path = DRAFT_MC_AR.path + if DRAFT_MC_AR.covers(rows): + workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) + out = None + if workspace is not None: + _record("m3_draft_mc_ar", rows, pushed) + out = m3_mc_allreduce.m3_mc_allreduce_norm( + x.contiguous(), + residual.contiguous(), + draft_model.w[norm_key], + draft_model.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + drafter_norm=True, + carveout=carveout, + prepushed=2 if pushed else 0, + ) + if out is not None: + _log_engaged(path, self.mapping) + if pushed: + taken(DRAFT_AR_PUSH_PATH, rows, "the projection pushed its rows (prepushed 2)") + elif DRAFT_AR_PUSH: + uncovered( + DRAFT_AR_PUSH_PATH, + "the boundary's projection did not push (the AR pushes x)", + ) + return out + if pushed: + raise RuntimeError( + "the drafter's rows were pushed into a boundary all-reduce that did not run" + ) + missed(path, rows, "the multicast all-reduce workspace or operand layout") + return draft_model.ar_norm(x, residual, norm_key) + + def _m3_arm_draft_push(self, draft_model: M3Eagle3Draft, key: str, rows: int) -> None: + """DRAFT_AR_PUSH: let projection ``key``'s next low-M call push its rows into the boundary all-reduce that + ``_m3_draft_ar_norm`` then runs (the multicast path only; shapes and a collectively created workspace decide + it, so the TP ranks agree).""" + if not (DRAFT_AR_PUSH and DRAFT_MC_AR.covers(rows)): + return + workspace = m3_mc_allreduce.get_workspace(self.mapping, draft_model.w[key].shape[0]) + if workspace is not None: + draft_model.arm_push(key, workspace, rows) + + def _m3_fused_loop_ok( + self, attn_metadata, spec_metadata, batch_size: int, draft_model: Optional[M3Eagle3Draft] + ) -> bool: + """Whether this step drafts with the fused loop: a decode-only step the fused verify or the seed prepared + (``_m3_first_gather_ids``) with the fused first draft inputs, a supported drafter, a (draft length, requests) + ``DRAFT_STEP`` covers, the prebuilt head and the exchange workspace. Host state only, so the TP ranks + agree.""" + if ( + self._m3_first_gather_ids is None + or not self._m3_fused_first_inputs + or draft_model is None + ): + return False + path = DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + uncovered( + path, "the drafter or its attention metadata is not one the fused loop serves" + ) + return False + draft_len = spec_metadata.runtime_draft_len + if not DRAFT_STEP.covers((draft_len, batch_size)): + uncovered(path, f"draft length {draft_len} at {batch_size} requests") + return False + head = self._m3_draft_head + if head is None or head.weight is not draft_model.lm_head.weight: + missed( + path, + batch_size, + "the drafter head's streamed copy is not built for its current weight (prebuild)", + ) + return False + if m3_draft_step_kernel.workspace(self.mapping, head.wf.device) is None: + capture_missed(path, None, "the draft-step exchange workspace") + return False + _log_engaged(path, self.mapping) + return True + + def _forward_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ): + """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop.""" + if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): + return self._m3_fused_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + return super()._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + + def _m3_fused_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model: M3Eagle3Draft, + draft_kv_cache_manager, + batch_size: int, + num_accepted_tokens, + ): + """Mirrors Eagle3OneModelWorker._forward_linear_draft_loop for a decode-only Eagle3 step, with each pass's + glue fused (see ``_m3_fused_loop_ok``): + + * the draft input -- the embedding lookup, the layer's input norm and hidden norm and their concat -- in one + launch (``m3_spec_tail.draft_input``; the second pass gathers the last accepted rows there); + * the drafter layer's attention and MLP (its SwiGLU on this target's C++ op), each followed by the all-reduce + with the residual add and the next RMSNorm (the post-attention norm, then the drafter's final norm): the + multicast boundary fed by the o_proj / down projection's push where covered (``_m3_draft_ar_norm``); + * one kernel for the BF16 head over the head rows (pass 0: the last accepted row of each request), the + cross-rank argmax and the bookkeeping the in-tree loop does between passes: the draft token, the next pass's + input id and position, the draft KV lengths (and pass 0's sequence lengths) (``m3_draft_step``). + + The host metadata transitions are the in-tree loop's. The drafts land in columns 1.. of the step's next-input + tokens (column 0 is the last accepted token), which ``_forward_impl`` returns as they lie. + """ + draft_len = spec_metadata.runtime_draft_len + table = draft_model.embed_table + dev = table.device + head = self._m3_draft_head + gather_ids = self._m3_first_gather_ids + ids = inputs["input_ids"] + hidden = inputs["hidden_states"] + pos = inputs["position_ids"] + if ids.dtype != torch.int32: + ids = ids.to(torch.int32) + if pos.dtype != torch.int32: + pos = pos.to(torch.int32) + num_acc = num_accepted_tokens[:batch_size] + if num_acc.dtype != torch.int32: + num_acc = num_acc.to(torch.int32) + next_new = self._m3_next_new + if next_new is None: + # Seeded from the in-tree acceptance: the last accepted token is the first-pass input at the gather row. + next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=dev) + next_new[:, 0] = ids[gather_ids] + draft_tokens = next_new[:, 1:] + w = draft_model.w + eps = draft_model.rms_eps + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) as md: + for i in range(draft_len): + x, residual = m3_spec_tail.draft_input( + ids, + table, + hidden, + gather_ids if i == 1 else None, + w["in_norm"], + eps, + w["hidden_norm"], + eps, + ) + if i == 0: + self._m3_join_pending() + rows = x.shape[0] + self._m3_arm_draft_push(draft_model, "o", rows) + h = draft_model.attention(x, md, reduce=False) + h, residual = self._m3_draft_ar_norm( + draft_model, h, residual, "post_norm", pushed=draft_model.take_pushed(h) + ) + self._m3_arm_draft_push(draft_model, "dn", rows) + h = draft_model.mlp(h, reduce=False, cpp_silu=True) + h, to_save = self._m3_draft_ar_norm( + draft_model, + h, + residual, + "final_norm", + carveout=DRAFT_MC_AR_CARVEOUT, + pushed=draft_model.take_pushed(h), + ) + if draft_model.return_hidden_post_norm: + to_save = h + ids_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + pos_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + if not m3_draft_step_kernel.draft_step( + head, + h, + batch_size, + gather_ids if i == 0 else None, + self.mapping, + draft_tokens, + i, + ids_next, + pos, + pos_next, + md.kv_lens_cuda, + num_acc if i == 0 else None, + md._seq_lens_cuda if i == 0 else None, + draft_len, + ): + raise RuntimeError( + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" + ) + ids, pos, hidden = ids_next, pos_next, to_save + if i == 0: + # The draft step set the device sequence lengths to 1 and rolled the KV lengths back. + md._seq_lens[:batch_size].fill_(1) + md.on_update() + has_kv_cache = md.kv_cache_manager is not None + if has_kv_cache and hasattr(md, "host_request_types"): + md.host_request_types[: md.num_contexts].fill_(1) + if has_kv_cache: + md.num_contexts = 0 + self._prepare_flash_mla_generation_layout(md, 0, batch_size) + md.update_for_spec_dec() + md.use_spec_decoding = False + else: + md.update_for_spec_dec() + self._m3_next_new_tokens = next_new + return draft_tokens + + +def _build_speculation(shell, model_config: ModelConfig): + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise ValueError( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + _set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return draft_config, draft, worker, use_separate + + +# -------------------------------------------------------------------------------------------------------------------- +# Runtime hooks and the shell +# -------------------------------------------------------------------------------------------------------------------- + + +def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: + """The step-loop hooks this target opts into. + + * ``pack_step_inputs``: a decode step's small input copies go out as one upload and one scatter at the graph + replay instead of a chain of copies the next graph waits on; the scatter is parallel (a program per copy), issued + by the decode graph's own prologue (this target's one-launch C++ op, ``m3_step_prologue``) from a window copied + on a side stream while the previous graph runs, so no operation is left between the step graphs and the replay + does not start with a host copy (its copy-engine -> compute turnaround was 10.3 us at c1). Sampled speculative + batches stage their Philox seed/offset uploads, and the accepted-draft counts, into the same pack; the sampler's + slot ids stay in a device buffer re-uploaded only when they change. + * ``pack_step_epoch``: the step epoch (step_inputs.step_epoch), one int64 device counter the decode graph's + prologue increments once per replay (program 0 of its first scatter launch, no extra kernel) and never resets. + The per-layer megakernel pieces tag their per-step device flags with it (contract v1: tag = epoch << 8 | layer). + * ``skip_unchanged_h2d``: steady decode re-uploads slot indices, query lengths and page tables that rarely change. + * ``h2d_changed_span``: a changed packed 2-D metadata upload stages only each row's changed column span. On a step + where a request crosses into a new page the MSA block table changed by one entry, but the whole padded table + (32 KB a request) was staged, overflowing the 16 KB in-graph window into an eager upload and scatter between the + step graphs (TP4 57-69 us at the ~1 in 19 block-crossing steps). + * ``kv_lens_update_once``: a speculative overlap step re-derives the MSA KV-length metadata once, after the + forward's KV-length correction (the first pass is rewritten in full by the second). + * ``spec_offsets_op``: the overlap scheduler's speculative step correction (position_ids += offsets, kv_lens += + offsets) as this target's one plain launch (``m3_spec_offsets``) instead of two ATen adds at the head of the + decode graph: the same int32 adds. + * ``loop_on_execution_stream``: the overlap loop's own stream work and the forward on one stream, so step N+1's + inputs wait on graph N in program order rather than through two cross-stream event hops. + * ``merge_tuned_tactics_across_ranks``: tensor-parallel peers capture the same GEMM tactics (rank-local profiling + noise otherwise splits them, and the slowest rank sets every step). + * ``skip_greedy_rng_upload``: all-greedy speculative batches never read the Philox seeds. + * ``spec_step_fold``: the spec sampler's slot stores and the engine's next-step input gathers; the decode forward + stores its outputs and stages the next step's inputs at its own tail (``M3SpecStepFold``). + * ``sampler_d2h_side_stream``: the spec sampler copies its per-slot store to the host on a side stream, overlapping + the next step's input staging. + * ``iteration_timing_one_event``: the iteration log times the device step with one event per boundary on a side + stream (the reported step times are the same). + * MSA: identity selections for short sequences (and short/long graph families), the one-launch KV-length update, + unchanged block-table copies skipped, pool views built once, block-id tables copied from the page-index views + (the stock builder makes a Python list of the whole padded row: on the Vera CPU 282 us per call at one request + and 561 us at two, against 8 and 13 us). + """ + return RuntimeHooks( + first_content_emit_iters=FIRST_CONTENT_EMIT_ITERS, + mxfp8_pinned_tactics=True, + late_compile_detect=True, + pack_step_inputs=True, + pack_parallel_scatter=True, + pack_in_graph=True, + pack_early_copy=True, + pack_step_epoch=True, + pack_prologue_op=m3_step_prologue, + skip_unchanged_h2d=True, + kv_lens_update_once=True, + spec_offsets_op=m3_spec_offsets, + h2d_changed_span=True, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + pack_spec_rng_uploads=True, + spec_small_h2d=True, + spec_step_fold=spec_fold, + sampler_d2h_side_stream=True, + iteration_timing_one_event=True, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=_msa_kv_update_kernel.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + block_ids_from_views=True, + ), + ) + + +class ModelingV2MiniMaxM3Nvfp4Sm107Tp8( + DecoderModelForCausalLM[MiniMaxM3Tp8ModelingV2, PretrainedConfig] +): + """MiniMax-M3 NVFP4 / sm_107 / tp8, text only, with one-model Eagle3 or no speculation. + + The shell owns the vocab-parallel BF16 ``lm_head`` and the logits processor (``DecoderModelForCausalLM``); with + speculation, the drafter (``M3Eagle3Draft``, its weights from the speculative checkpoint through + ``load_draft_weights``) and the worker (``M3Eagle3Worker``) that verifies and drafts each step. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + assert torch.cuda.get_device_capability() == _SM, ( + f"this target is built for sm_{_SM[0]}{_SM[1]}, the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + # Under modeling_v2='require' a covered decode shape that cannot take its fast path raises. + set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + model_config = _text_model_config(model_config) + cfg = model_config.pretrained_config + spec_config = getattr(model_config, "spec_config", None) + if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: + # The one-model Eagle3 layer has no NVFP4 trtllm-gen decode kernel: keep its modules and the shared + # draft cache on FP8 while the target stays NVFP4. + model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 + super().__init__( + MiniMaxM3Tp8ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self.spec_config = spec_config + self.draft_model: Optional[M3Eagle3Draft] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[M3Eagle3Worker] = None + self.use_separate_draft_kv_cache = False + self.layer_idx = FINAL_HIDDEN_LAYER_ID + self._spec_fold: Optional[M3SpecStepFold] = None + if spec_config is not None: + ( + self.draft_config, + self.draft_model, + self.spec_worker, + self.use_separate_draft_kv_cache, + ) = _build_speculation(self, model_config) + self._spec_fold = M3SpecStepFold() + self.runtime_hooks = _runtime_hooks(self._spec_fold) + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights: Dict, weight_mapper=None) -> None: + """The drafter's checkpoint (``speculative_model``), loaded by the engine after the target's.""" + _weights.DRAFT_WEIGHTS.load(self.draft_model, weights) + + def set_guided_decoder(self, guided_decoder) -> bool: + """Mirrors ``SpecDecOneEngineForCausalLM.set_guided_decoder``: guided decoding with speculation is installed + on the worker.""" + if hasattr(self.spec_worker, "set_guided_decoder"): + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + def cache_derived_state(self) -> None: + """Build the targets and every decode kernel once the weights are real, before the executor's warmup; with + speculation also the drafter's full-context decode, the fused verify's coded target head, the die-placed + target head copy, the fused draft loop's head copy and the drafter linears' low-M kernels.""" + super().cache_derived_state() + self.model.post_load_weights() + if self.draft_model is not None and DRAFTER_FULLCTX.keys: + path = DRAFTER_FULLCTX.path + recipes = {k: FULLCTX_DECODE_RECIPES[k] for k in DRAFTER_FULLCTX.keys} + installed = m3_drafter_attn.install( + self.draft_model, + recipes, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and m3_fullctx_decode_kernel.has_own_reducer(), + get_env_enable_pdl(), + fused_rope_mode=( + DRAFTER_ROPE_MODE + if DRAFTER_ROPE_FUSED + and hasattr(torch.ops.trtllm, "minimax_m3_sm107tp8_drafter_rope_fp8_kv") + else None + ), + on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), + on_taken=lambda rows: taken(path, rows), + on_uncovered=lambda reason: uncovered(path, reason), + on_missed=lambda rows, reason: missed(path, rows, reason), + ) + logger.info(f"[M3 prebuild] {path} installed on {installed} drafter attention layer(s)") + if self.spec_worker is not None: + self.spec_worker.prebuild(self.lm_head, self.draft_model) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + """One engine step: the target forward, then the logits (or, with speculation, the worker's verify and + draft, and on a decode step the spec fold's store of its outputs). Text only.""" + multimodal_params = kwargs.pop("multimodal_params", None) or [] + if any(p is not None and getattr(p, "multimodal_data", None) for p in multimodal_params): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + worker = self.spec_worker + capture = _plain_capture + capture_target = None + if worker is not None: + capture = worker.capture_hidden_states + capture_target = worker.capture_target + if spec_metadata is not None: + worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + capture=capture, + capture_target=capture_target, + **kwargs, + ) + if getattr(attn_metadata, "csa2_remote_tail_mode", None) == "source": + # A context worker stopped at the encoder/decoder boundary: a disposable token for the context-only + # completion, and no draft worker. + return self.logits_processor.forward(hidden_states, self.lm_head, attn_metadata, False) + if spec_metadata is not None and spec_metadata.is_layer_capture(FINAL_HIDDEN_LAYER_ID): + capture(spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + if worker is None: + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], self.lm_head, attn_metadata, True + ) + spec_input_ids = input_ids if input_ids is not None else kwargs.get("orig_input_ids") + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if spec_input_ids is not None: + spec_input_ids = spec_input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens) + outputs = worker( + input_ids=spec_input_ids, + position_ids=spec_position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + spec_fold = self._spec_fold + if ( + spec_fold is not None + and spec_fold.in_forward + and attn_metadata.num_contexts == 0 + and isinstance(outputs, dict) + ): + # The decode forward stores its outputs into the spec sampler's store and stages the next step's inputs + # at its tail (one launch inside the step graph). + spec_fold.store_in_forward(outputs) + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/__init__.py new file mode 100644 index 000000000000..b8c7426e3933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Activation ops of this target: the gated MLPs' fused SwiGLU.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..e7ceefc5bdab --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SiluAndMul(OpWrapper): + """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with + ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain + SwiGLU ``silu(gate) * up``. + + The MXFP8 dense MLPs (layers 0-2, intermediate 6144 per rank) and the MoE's shared expert (intermediate 1536 per + rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 9216 per + rank) runs it plain. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=384, oai=True)), + Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=384, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=1536, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=1536, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=2304, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=2304, oai=False), + ), + ) + + note = """ + The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two + differ only by the fast sigmoid's last bits (the test's band). + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + gate, up = x.float().chunk(2, dim=-1) + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + a = 1.0 if alpha is None else alpha + b = 0.0 if beta is None else beta + return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" + + +m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/__init__.py new file mode 100644 index 000000000000..a7f7fee7a234 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention ops of this target: MSA block selection, the sparse decode and the KV-length update.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..789ae79db88e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per +request count (modeling's ``DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 16 +KV_HEADS = 1 +HEAD_DIM = 128 +#: The split this target hands trtllm-gen at one request (modeling's DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model +#: maximum the in-tree step passes otherwise. +CAPPED = 16384 +MODEL_MAX = 1 << 20 + + +def dense_attention_reference( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first + ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's + tokens in order). Returns fp32 [total_q, heads, head_dim].""" + batch, kv_heads, max_len, _ = keys.shape + heads = q.shape[1] + group = heads // kv_heads + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + pos = torch.arange(max_len, device=q.device) + with true_fp32_matmul(): + for b in range(batch): + full = int(seq_lens[b]) + rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] + limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) + valid = pos[None, :] < limit[:, None] # [qlen, max_len] + for h in range(kv_heads): + qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] + logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + probs = torch.softmax(logits, dim=-1) + out[ + b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group + ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) + return out + + +def gather_paged_kv( + kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool + (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + page = pool.shape[2] + pages = (int(seq_lens.max()) + page - 1) // page + slots = block_table[:, :pages].long() + k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] + v = pool[slots * subpages + 1] + batch, _, kv_heads, _, d = k.shape + + def order(t): + return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) + + return order(k), order(v) + + +class M3DenseDecode(OpWrapper): + """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len + ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. + + Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into + ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: + every CTA still covers its share of each request's actual length, which may exceed it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="two requests keep the model maximum (no cap at two requests)", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), + ), + ) + + note = """ + The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped + one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV + band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + minimax_m3_trtllm_gen_dense_decode( + q, + kv_cache_manager, + layer_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + max_seq_len=max_seq_len, + max_num_requests=max_num_requests, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). + + ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + if keys.dtype == torch.float8_e4m3fn: + q = fp8_e4m3_saturate(q) + return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> None: + # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. + assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" + # Each verify row attends a causal prefix of at least one key. + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + + +m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..c2b23aa18750 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention on the TRTLLM backend: RoPE, KV-cache append and causal attention.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class M3DraftAttention(OpWrapper): + """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the + backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's + NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` + ``num_tokens``); returns [num_tokens, q heads x 128] bf16. + + ``backend`` (``TrtllmAttention``) and ``metadata`` are the runtime's objects for the drafter layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="first draft pass of a 2-request verify (2 x 4 rows)", spec=dict(reqs=2, qlen=4)), + Cell(why="a later draft pass, one row per request", spec=dict(reqs=2, qlen=1)), + Cell(why="a context chunk's draft rows", spec=dict(reqs=1, context=512)), + ) + + note = """ + `reference` is exact (float64) causal softmax attention: Q rotated by the backend's NeoX RoPE at each token's + position, over the drafter layer's K / V rows read back from its FP8 cache (dequantized by the backend's + ``kv_scale_quant_orig``). What remains between the two is the FMHA's own arithmetic (Q's rounding to the + kernel's input precision, its P rounding and accumulation order), which the test bounds. + """ + + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors grouped KV heads (query head h reads KV head h // (heads / kv heads); two per rank at tp2) and + un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + # [kv heads, tokens, dim]: each KV head's rows in token order across the pages. + keys = ( + pool[pages, 0] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + values = ( + pool[pages, 1] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + keys = keys.repeat_interleave(group, dim=0) # [heads, tokens, dim] + values = values.repeat_interleave(group, dim=0) + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] + rows = q[first : first + n] + x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,hkd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.einsum( + "thk,hkd->thd", torch.softmax(logits, dim=-1), values + ) + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) + + +m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..2fecad0e9afa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q +out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3Fp8IndexerQkNormRope(OpWrapper): + """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and + writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, + then RoPE on its leading ``rotary_dim`` channels (NeoX). + + The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; + the fused layers' producer writes the index cache itself. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, two index heads per rank", spec=dict(m=1, heads=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, heads=2)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=2)), + ) + + note = """ + `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 + before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of + values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. + """ + + def __call__( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + index_k_cache, + out_cache_loc, + num_heads_q, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + position_ids, + ) + + def reference( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, + then the E4M3 cast.""" + out = fused_qk_norm_rope( + qk.clone(), + num_heads_q, + 1, + 0, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) + return ( + q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), + k.to(torch.float8_e4m3fn), + ) + + +m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py new file mode 100644 index 000000000000..a5439efc660b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' full-context decode attention at TP8 on Rubin (``kernels/m3_fullctx_decode``): PrimTS split-KV +decode at the recipe modeling's ``FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) + +from ...kernels import m3_fullctx_decode as kernel +from .m3_dense_decode import dense_attention_reference, gather_paged_kv + +__all__ = ["m3_fullctx_decode"] + +HEADS = kernel.HEADS +KV_HEADS = kernel.KV_HEADS +HEAD_DIM = kernel.HEAD_DIM +#: The plan capacity (the per-request K/V length bound PrimTS sizes its policy for): dispatch's +#: FULLCTX_DECODE_MAX_KV_LEN, which the forward's prebuilt plans use. +MAX_KV_LEN = 1 << 21 + + +class M3FullCtxDecode(OpWrapper): + """Decode attention of every query row over its request's full causal context in the paged FP8 pool, written + into ``output`` (BF16 [T, heads, head_dim]) in place. + + ``recipe`` is ``modeling's FULLCTX_DECODE_RECIPES[(qlen, requests)]``: the PrimTS Q tile and the split-KV fanout. + ``own_reducer`` merges the splits on the BF16 op (modeling's ``FULLCTX_DECODE_OWN_REDUCER``) rather than PrimTS's + FP16 reducer and a cast. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over an AgentX-length (165k) context", + spec=dict(kv_lens=(165525,), qlen=4), + ), + Cell( + why="one request, one draft-pass row, over a 165k context", + spec=dict(kv_lens=(165522,), qlen=1), + ), + Cell( + why="two requests verifying 1 + 3 rows at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165525, 120011), qlen=4), + ), + Cell( + why="two requests, one row each, at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165522, 120008), qlen=1), + ), + Cell( + why="one request verifying 1 + 3 rows over a short (4.1k) context: most splits are empty", + spec=dict(kv_lens=(4104,), qlen=4), + ), + ) + + note = """ + Same operands and arithmetic as the stock trtllm-gen call (E4M3 q, E4M3 K/V, P quantized to E4M3 per split, + fp32 accumulation): the split differs, so the result differs from stock as a stock split change does (the P + quantization is per split). Deterministic, and the same bits eager and in a CUDA graph. Held to the fp32 + reference within the FP8-KV band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + shape = (decode_query_len, int(seq_lens.shape[0])) + kernel.prebuild(q.device, {shape: recipe}, max_kv_len, own_reducer, get_env_enable_pdl()) + plan = kernel.plan_for(q.device, *shape, own_reducer) + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + k_pages, v_pages = kernel.kv_page_views(pool, subpages) + plan.run(fp8_e4m3_saturate(q), k_pages, v_pages, block_table, seq_lens, sm_scale, output) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V, q rounded to E4M3 first as the kernel runs it.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + return dense_attention_reference( + fp8_e4m3_saturate(q), keys, values, seq_lens, decode_query_len, sm_scale + ) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> None: + pool, _ = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + assert pool.dtype == torch.float8_e4m3fn, "an FP8 E4M3 pool" + assert tuple(pool.shape[1:]) == (KV_HEADS, kernel.PAGE, HEAD_DIM), ( + "the TP8 dense-layer pool geometry" + ) + assert q.shape[1:] == (HEADS, HEAD_DIM), "8 query heads of 128 per rank" + assert int(q.shape[0]) == int(seq_lens.shape[0]) * decode_query_len, ( + "q covers exactly the batch" + ) + if torch.cuda.is_current_stream_capturing(): + return # the length checks read the device: eager calls only + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + assert bool((seq_lens <= max_kv_len).all()), "a request longer than the plan capacity" + + +m3_fullctx_decode = M3FullCtxDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..674b16e7744f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``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 + + +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 kernel at every request count it covers) issues the Q and first K tiles right after the + grid wait from a pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise + every older block of the CTA is prefetched. Neither, nor ``split_k``, changes a score. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request, 1 + 3 rows over 165521 tokens (AgentX ISL), 1 index head (TP8): 1536 CTAs, early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most of the 1536 CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request of 400000 tokens: 3125 blocks, so the 1536 CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request decoding one row over 65536 tokens", + spec=dict(seq_lens=(65536,), qlen=1, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request, 1 + 3 rows over 165521 tokens, 1 index head (TP8): 1272 CTAs (one wave), early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="one request of 129 tokens at 1272 CTAs: the early tiles of most CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="one request of 400000 tokens at 1272 CTAs: 3125 blocks, so the CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="two requests verifying 4 rows each over 165521 and 20000 tokens: 768 CTAs each, early issue", + spec=dict(seq_lens=(165521, 20000), qlen=4, heads=1, split_k=768, early=True), + ), + Cell( + why="two requests, one of 300 tokens and one of 262144 (2048 blocks: the 768 CTAs loop)", + spec=dict(seq_lens=(300, 262144), qlen=4, heads=1, split_k=768, early=True), + ), + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..f1a7f210a590 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MiniMax-M3 MSA attention backend's step: KV write, sparse top-k block selection, FMHA.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + +from .m3_sparse_decode_lc8 import dequant_nvfp4_pages + + +class M3MsaAttention(OpWrapper): + """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` + [T, heads * 128] (bf16, the step's live rows). + + * ``k`` / ``v`` given (the dense layers): written into the layer's paged main cache first (NVFP4 with + ``kv_scale_orig_quant``, else FP8 / BF16 as the pool is). None: a producer already wrote the caches. + * ``sparse``: the indexer scores the layer's index-K cache (already written) with ``idx_q`` and publishes the + top-k blocks; the FMHA attends those blocks. Otherwise the FMHA attends the whole page table. + + The FMHA dispatches on the step's metadata (context rows: the MSA prefill FMHA; generation rows: the MSA decode + FMHA, the dense layers' trtllm-gen decode at their capped KV split). ``backend`` and ``metadata`` are the runtime's + objects for this layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dense layer, a prefill chunk, FP8 pool", + spec=dict(layer="dense", phase="context", tokens=1024), + ), + Cell( + why="dense layer, a decode step of 2 requests", + spec=dict(layer="dense", phase="generation", reqs=2), + ), + Cell( + why="sparse layer, a prefill chunk beyond 16 blocks, NVFP4 pool", + spec=dict(layer="sparse", phase="context", tokens=4096), + ), + Cell( + why="sparse layer, a mixed step (context + generation rows)", + spec=dict(layer="sparse", phase="mixed", tokens=512, reqs=2), + ), + ) + + note = """ + `reference` is exact (float64) softmax attention over the cache rows the step attends, read back from the paged + cache the call wrote (or a producer wrote before it) and dequantized as the FMHA dequantizes them: FP8 rows as + E4M3, NVFP4 rows as E2M1 x E4M3 block scale x the per-tensor ``kv_scale_quant_orig``. A sparse layer's top-k blocks + are re-selected from ``idx_q`` and the index-K cache with the backend's rule (per-block max over the causal + keys, the query's own block forced in, ties to the lower block). On an FP8 pool the FMHA takes Q as a plain E4M3 + cast, which the reference applies too. What remains between the two is the FMHA's own arithmetic (its P + rounding and accumulation order), which the test bounds. + """ + + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors grouped KV heads (query + head h reads KV head h // (heads / kv heads); two per rank at tp2), each KV head selecting its own top-k + blocks from its own index heads against the rank's single index-K head. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + assert config.num_index_heads % kv_heads == 0 + idx_group = config.num_index_heads // kv_heads + + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + qpos = cached[row] + torch.arange(n, device=q.device) + causal = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + for g in range(kv_heads): + if nvfp4: + keys = dequant_nvfp4_pages( + pool[:, 0].view(torch.uint8), + scale_pool[:, 0], + pages, + quad_scales=False, + kv_head=g, + ) + values = dequant_nvfp4_pages( + pool[:, 1].view(torch.uint8), + scale_pool[:, 1], + pages, + quad_scales=True, + kv_head=g, + ) + keys, values = keys * k_global, values * v_global + else: + keys = pool[pages, 0, g].to(torch.float64) + values = pool[pages, 1, g].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + allowed = causal + if sparse: + # One index-K head per rank, shared by every index head (the cache is [pages, 1, page, dim]). + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = _selected_blocks( + idx64[first : first + n, g * idx_group : (g + 1) * idx_group], + idx_keys, + causal, + qpos, + page, + config, + ) + allowed = causal & chosen.repeat_interleave(page, dim=1)[:, :length] + h0, h1 = g * group, (g + 1) * group + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = ( + torch.einsum("thd,kd->thk", q64[first + lo : first + hi, h0:h1], keys) + * sm_scale + ) + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi, h0:h1] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + (ids >= n_valid - config.local_blocks) & (ids < n_valid) + ) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen + + +m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py new file mode 100644 index 000000000000..ff34b8eac631 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer, with its constant loads staged +before the PDL grid-dependency wait and an early dependents trigger (``kernels/minimaxM3/producerPw.cu``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +__all__ = ["m3_producer"] + +HEAD_DIM = 128 +ROTARY_DIM = 64 +PAGE = 128 +#: Query heads per rank at TP8 (MiniMax-M3: 64 query heads; 1 KV head and 1 index head per rank). +HEADS_Q = 8 +#: NVFP4 block: 16 values share one E4M3 scale. +FP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _norm_rope( + x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(FP4_BLOCK) + + +def scale_index(within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return (within_page // 4) * 32 + cols * 4 + within_page % 4 + + +class M3Producer(OpWrapper): + """The NVFP4 QKV + indexer producer of a sparse layer: ``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert`` for one + KV head, one index head, head dim 128 and rotary dim 64, with the same arguments and outputs bit for bit. + + ``packed`` [T, (heads_q + 4) * 128] BF16 holds each token's q heads, k, v, index-q and index-k. Every head but v + gets the Gemma RMSNorm with its weight and NeoX RoPE on its first 64 dims (``rotary_cos_sin[position]``). Returns + (q FP8 [T, heads_q, 128], index-q FP8 [T, 1, 128]); writes the token's slot ``out_cache_loc[t]`` (skipped when + negative or past the pool): k and v as NVFP4 (16-value blocks, E4M3 scales against ``kv_quant_scale[1]`` / + ``[2]``, the V scales in the 4-token swizzle) and index-k as FP8. ``early`` triggers the dependent kernel on + entry (its outputs are read only after the dependent's own grid wait). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token (c1 without speculation)", spec=dict(tokens=1, scale=1.0)), + Cell(why="a c1 verify: 4 tokens", spec=dict(tokens=4, scale=3.0)), + Cell(why="a c4 verify: 16 tokens, small activations", spec=dict(tokens=16, scale=1e-3)), + Cell( + why="a c8 verify: 32 tokens, large activations, a skipped and an out-of-pool slot", + spec=dict(tokens=32, scale=40.0, bad_slots=True), + ), + Cell(why="the widest decode step: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="a prefill chunk of 1000 tokens, early trigger off", + spec=dict(tokens=1000, scale=1.0, early=False), + ), + ) + + note = """ + The kernel is the native producer's arithmetic, compiled with the library's --use_fast_math, so it is held to that + kernel bit for bit on every output and cache byte. The torch reference states the math in fp32: the FP8 outputs are + held to it within one E4M3 step, and the NVFP4 K/V, decoded, within one E2M1 step of their block scale. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if prewait_step_inputs: + # This target's copy (producerPw8): the same outputs bit for bit, with the position ids, slots and RoPE + # rows read before its grid wait (none of them may be written by a PDL predecessor). + return tuple( + torch.ops.trtllm.minimax_m3_sm107_tp8_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + True, + ) + ) + q, index_q = torch.ops.trtllm.minimax_m3_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + ) + return q, index_q + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", + "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change + them.""" + t = packed.shape[0] + x = packed.float().view(t, heads_q + 4, HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = rope[:, 0], rope[:, 1] + q = torch.stack( + [_norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + ) + return { + "q": q, + "k": _norm_rope(x[:, heads_q], k_weight, cos, sin, eps), + "v": x[:, heads_q + 1], + "index_q": _norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), + "index_k": _norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), + } + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> None: + t = packed.shape[0] + # The RoPE row is read at the position with no bound. + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + # Two tokens writing one slot race. + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" + + +m3_producer = M3Producer() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..88180c2f1a1f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3QkNormRopeFp8(OpWrapper): + """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded + once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place + entry. + + The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps + of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp8 heads (8 q, 1 kv)", spec=dict(m=1, q=8, kv=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=8, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=8, kv=1)), + ) + + note = """ + `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from + fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. + """ + + def __call__( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + + def reference( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" + out = fused_qk_norm_rope( + qkv.clone(), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return out.to(torch.float8_e4m3fn) + + +m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..c6f4be426ba3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into +the paged caches.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3QkvInsert(OpWrapper): + """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the + partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and + index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. + + ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool + [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 + [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp8 + heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), + Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), + ) + + note = """ + The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` + returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + common = ( + num_heads, + num_kv_heads, + num_index_heads, + head_dim, + rotary_dim, + eps, + q_norm, + k_norm, + index_q_norm, + index_k_norm, + rotary_cos_sin, + position_ids, + ) + if kv_scale is not None: + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common + ) + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, index_k_cache, out_cache_loc, *common + ) + + def reference( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of + Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as + packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" + m = packed.shape[0] + sizes = [ + num_heads * head_dim, + num_kv_heads * head_dim, + num_kv_heads * head_dim, + num_index_heads * head_dim, + head_dim, + ] + q, k, v, iq, ik = packed.float().split(sizes, dim=-1) + cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) + half = rotary_dim // 2 + + def norm_rope(x, w, heads): + x = x.reshape(m, heads, head_dim) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) + x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] + c, s = cos.unsqueeze(1), sin.unsqueeze(1) + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) + + return { + "q": norm_rope(q, q_norm, num_heads), + "index_q": norm_rope(iq, index_q_norm, num_index_heads), + "k": norm_rope(k, k_norm, num_kv_heads), + "v": v.reshape(m, num_kv_heads, head_dim), + "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], + } + + +m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..76841cca732d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SelectBlocks(OpWrapper): + """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), + with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the + row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` + backs the result head-major so ``permute(1, 0, 2)`` is contiguous. + + The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="eight decode rows of 17..128 blocks", + spec=dict(total_q=8, n_blocks=128, mixed=True), + ), + Cell( + why="one decode row at the 128-block edge", + spec=dict(total_q=1, n_blocks=128, mixed=False), + ), + Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), + ) + + note = """ + Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. + """ + + def __call__( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + return torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output + ) + + def reference( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + """Plain torch selection; ``head_major_output`` changes the backing only.""" + heads, n_blocks, total_q = max_score_kv.shape + out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) + blocks = torch.arange(n_blocks, device=max_score_kv.device) + for t in range(total_q): + nv = int(n_valid_blocks[t]) + for h in range(heads): + s = max_score_kv[h, :, t].clone().double() + s[blocks >= nv] = float("-inf") + forced = (blocks < min(init_blocks, nv)) | ( + (blocks >= nv - local_blocks) & (blocks < nv) + ) + s[forced] = float("inf") + order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] + ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) + out[t, h, : ids.numel()] = ids + return out + + +m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..9a92050c6abb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). +MAX_BLOCKS = 9216 +#: Blocks the selection keeps per (query, KV head). +TOPK = 16 +#: Scores the forcing substitutes: initial blocks rank above local ones, and both above any finite score. +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 + +# M3's checkpoint at TP8 (one KV head per rank: 4 KV heads over 8 ranks, each group on a rank pair): no forced initial +# blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=1, init_blocks=0, local_blocks=1) + + +class M3SelectLong(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query's top-16 blocks. + + ``scores`` is the per-KV-head block maxima [num_kv_heads, num_blocks, + total_q] the CuTe DSL index scorer wrote (amax over the KV head's index + heads). Per (query, KV head) row the op keeps the ``n_valid_blocks[query]`` + leading blocks (clamped to ``[0, num_blocks]``), forces the last + ``local_blocks`` of them (counted from the *unclamped* valid count) to + ``LOCAL_SCORE`` and the first ``init_blocks`` to ``INIT_SCORE`` -- local wins + where both apply -- and selects the 16 largest, ties to the smaller block. + A pick whose forced score is -inf is emitted as -1. The ids are written + ascending with a -1 tail. + + Same contract as ``torch.ops.trtllm.minimax_m3_select_blocks`` (the in-tree + ``select_blocks_from_maxscore``) and bit-identical to it, except on a row + where more than 2048 valid blocks tie at the 16th-largest score: there the + in-tree op fills the tied slots in atomic order (its staging holds 2048 + candidates), while this op breaks the tie to the smallest blocks as the + contract says. Up to 2048 tied blocks both take the smallest. The test checks + both against ``reference``. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the TP2 engine's 5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + note = """ + Certified at one KV head per rank (TP8, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`). The kernel takes any head count and + forcing; other forcing is not driven here. + + `scores` may be strided (the op reads it by stride); the target hands over + the scorer's contiguous view, which is what the cells drive. Rows of at + most 128 blocks are within the kernel's range, but the target keeps those on + the in-tree op, so they are not certified here. + + Ordering is by the descending-score key `orderedKey`: a positive NaN ranks + above +inf, +0 above -0, and a negative NaN below -inf -- which, unlike + -inf, is emitted as a block id rather than -1. The cells drive positive + NaN only. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sm107_tp2_select_long( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and + ``trigger_phase`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() # [total_q, num_blocks] + forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where( + in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max + ) + k = min(TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) + ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> None: + """The inputs this op takes and answers wrongly. + + Negative forcing counts are accepted and push the forced ranges past the + row instead of failing; the in-tree op this one stands in for refuses + them, so a caller relying on this one would silently diverge from it. + An ``out`` overlapping ``scores`` is overwritten while it is still read. + """ + assert init_blocks >= 0 and local_blocks >= 0, ( + f"init_blocks ({init_blocks}) and local_blocks ({local_blocks}) must be non-negative; the " + "op accepts negative counts and selects as if the forced ranges were elsewhere" + ) + storage = scores.untyped_storage() + assert not (storage.data_ptr() <= out.data_ptr() < storage.data_ptr() + storage.nbytes()), ( + "out must not alias scores: the op writes it while scores are still read" + ) + + +m3_select_long = M3SelectLong() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py new file mode 100644 index 000000000000..cd76a4e6d5d2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_select_long import _M3, MAX_BLOCKS, M3SelectLong + + +class M3SelectLongDsl(OpWrapper): + """``m3_select_long``'s contract, cells and reference, run by ``kernels/m3_select_long_dsl``: the same selection + bit for bit as the 1024-thread ``minimax_m3_sm107_tp2_select_long`` (also on rows where more than 2048 blocks tie + at the 16th-largest score), from one 384-thread CTA per row. Returns ``out``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the TP2 engine's 5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + """``m3_select_long``'s reference: the selection is the same bit for bit.""" + return M3SelectLong.reference( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> None: + """``m3_select_long``'s domain.""" + M3SelectLong.is_valid( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + + note = """ + Certified at one KV head per rank (TP8, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`), and bit for bit against the + 1024-thread `m3_select_long` on every cell. The kernel takes any head count + and forcing; other forcing is not driven here. + + The kernel is compiled per (PDL launch, early trigger, trigger phase) with + symbolic shapes and strides; the target compiles its keys before warmup + (`kernels/m3_select_long_dsl.prebuild`), and a call outside them compiles + on first use. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + from ...kernels.m3_select_long_dsl import select_long_dsl + + if scores.shape[1] > MAX_BLOCKS: + raise RuntimeError( + f"m3_select_long_dsl supports at most {MAX_BLOCKS} blocks, got {scores.shape[1]}" + ) + select_long_dsl( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + +m3_select_long_dsl = M3SelectLongDsl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..6525a3341b15 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This target's block-sparse decode attention (Triton, with the one-launch long-context route).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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 (one KV head) with its block and per-tensor + scales, and BF16 ``output`` [T, 8, 128]. Returns ``output``. + + Which kernel runs is decided inside, per step: a selected (non-identity) table on the TP8 geometry goes to the + one-launch ``m3_sparse_decode_lc8`` kernel while modeling's ``SPARSE_DECODE_LC8`` covers it; an identity table + (every request within 16 blocks, staged by the metadata), or a selected table with that range killed, runs the + Triton per-row producer (the multi-query producer for identity steps of at least 17 rows), then the LSE merge. + Every route also publishes the MXFP8 o_proj hand-off to ``m3_prequant.SLOT``, bit for bit + ``torch.ops.trtllm.mxfp8_quantize(output, True)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="identity selection, decode without speculation at batch 8: the Triton producer + merge", + spec=dict( + seq_lens=(1, 128, 129, 700, 2048, 1500, 64, 0), decode_query_len=1, identity=True + ), + ), + Cell( + why="identity selection, Eagle3 verify at one request (T=4): the Triton producer + merge", + spec=dict(seq_lens=(1900,), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 2 (T=8), a padding request", + spec=dict(seq_lens=(2048, 0), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 8 (T=32): the multi-query producer + merge", + spec=dict( + seq_lens=(1900, 2048, 700, 1500, 1200, 2000, 300, 1024), + decode_query_len=4, + identity=True, + ), + ), + Cell( + why="selected blocks, Eagle3 verify at one request (T=4): the one-launch long-context kernel", + spec=dict(seq_lens=(40000,), decode_query_len=4, identity=False), + ), + Cell( + why="selected blocks, Eagle3 verify at batch 2 (T=8): the one-launch long-context kernel", + spec=dict(seq_lens=(9000, 40000), decode_query_len=4, identity=False), + ), + ) + + note = """ + Certified only as the target calls it: one KV head of 8 query heads, head_dim 128, FP8 q, packed NVFP4 + cache, BF16 output, `num_topk_chunks` left to the dispatcher. The function also takes BF16 / FP8 caches, other + head geometries and an explicit split-K factor (the in-tree tests drive those); none is certified here. + + The MXFP8 hand-off of the long-context route is certified on its own entry (`m3_sparse_decode_lc8`); the Triton + merge's (`_store_mxfp8_head`) is held to `torch.ops.trtllm.mxfp8_quantize` by this entry's test. + """ + + def __call__( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + kv_scale=kv_scale, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + num_topk_chunks=num_topk_chunks, + identity_blocks=identity_blocks, + ) + return output + + def reference( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + """``output`` as the call is supposed to write it: float64 block-sparse attention over the + dequantized NVFP4 blocks (``sparse_attention_reference``), rounded once to ``output``'s dtype. + + Only the packed NVFP4 cache is mirrored -- the target's only configuration -- one KV head at a time: KV head + ``h`` serves query heads ``[h * group, (h + 1) * group)`` with its own top-k row. ``kv_scale``, + ``num_topk_chunks`` and ``identity_blocks`` do not change what is computed. + """ + assert k_block_scale is not None, "mirrors the NVFP4 cache only" + kv_heads = k_paged.shape[1] + assert kv_heads == 1, "mirrors the TP8 cache (one KV head per rank) only" + group = q.shape[1] // kv_heads + heads = [] + for h in range(kv_heads): + heads.append( + sparse_attention_reference( + q[:, h * group : (h + 1) * group], + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[h], + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + ) + return torch.cat(heads, dim=1).to(output.dtype) + + def is_valid( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) + + +m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..a845ac68f257 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch long-context sparse decode over the selected NVFP4 blocks at 8 query heads (16-CTA clusters). + +Also the plain-torch block-sparse attention and MXFP8 hand-off mirrors the +Triton sparse decode entry (``m3_sparse_decode``) reuses: both compute the same +thing over the same cache. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The TP8 geometry the kernel is compiled for: 8 query heads on one KV head. +NUM_HEADS = 8 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +#: NVFP4 elements per E4M3 block scale. +SF_VEC = 16 +#: MXFP8 elements per UE8M0 scale, and the o_proj width the hand-off is laid out for. +MX_BLOCK = 32 +WIDTH = NUM_HEADS * HEAD_DIM + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, + block_scale: torch.Tensor, + pages: torch.Tensor, + quad_scales: bool, + kv_head: int = 0, +) -> torch.Tensor: + """Dequantize KV head ``kv_head`` of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, kv_head].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, kv_head].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """Block-sparse decode attention over one KV head of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages(k, k_block_scale, pages, quad_scales=False).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True).reshape(-1, HEAD_DIM) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, + sf [num_rows_alloc * WIDTH // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) + num_sf_k_tiles = WIDTH // 128 + sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +class M3SparseDecodeLc8(OpWrapper): + """Sparse decode attention, split-K merge and MXFP8 hand-off of one TP8 decode step in one launch. + + One cluster of 16 CTAs per query token: CTA ``r`` attends the token's + ``r``-th selected block (``topk[tok, r]``) for all 8 heads (transposed MMAs, + the heads on the n dimension), and the cluster merges the 16 partials over + distributed shared memory into the BF16 row ``out[tok]`` [1024]. With ``q8`` / + ``sf`` it also writes the row's MXFP8 copy and swizzled UE8M0 scales for the + o_proj GEMM. All operands are byte views (E4M3 ``q``, packed E2M1 pages, E4M3 + block scales). The kernel releases its PDL dependents right after its grid + wait. Returns ``out``. + + ``prewait`` stages the request length, the query rows (widened to FP16) and + the request's block-table row (up to 8192 entries, longer rows are read from + global memory) before the grid wait; the arithmetic is the same, so both + variants write the same bits. + + Per block: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16 + (exact), QK and PV on FP16 tensor cores with FP32 accumulation, the softmax + probabilities rounded to FP16 for PV, log2-domain softmax with ``ex2.approx``. + ``linear_softmax`` picks the denominator recurrence (the two agree to + rounding); the target picks it per batch shape. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation, one request: T=1, log-domain recurrence, MXFP8 hand-off, prewait", + spec=dict( + seq_lens=(9000,), decode_query_len=1, linear_softmax=False, mxfp8=True, prewait=True + ), + ), + Cell( + why="Eagle3 verify at one request (the TP8 operating point): T=4, log-domain recurrence, prewait", + spec=dict( + seq_lens=(20000,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="Eagle3 verify at two requests: T=8 on the linear denominator, prewait at its row limit", + spec=dict( + seq_lens=(5000, 12345), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="a 1M-token request: a block-table row longer than the staged 8192 entries, read from global", + spec=dict( + seq_lens=(1048576,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="Eagle3 verify at batch 8 (T=32): short requests with -1 tails and zero-length padding rows", + spec=dict( + seq_lens=(16384, 2049, 130, 0, 7777, 4, 30000, 0), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="decode without speculation at batch 8: T=8 rows of distinct requests, staged after the wait", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="no hand-off buffers (the variant the call site takes past 128 rows), driven at T=8", + spec=dict( + seq_lens=(5000, 12345), + decode_query_len=4, + linear_softmax=True, + mxfp8=False, + prewait=False, + ), + ), + ) + + note = """ + Certified: `linear_softmax` x hand-off on / off x `prewait`, the variants + the cells drive. The hand-off cells are certified bit-exactly against + `mxfp8_handoff_reference` applied to the op's own BF16 rows; the BF16 rows + themselves within the band the test derives, and the two `prewait` variants + against each other bit for bit. + + `topk` is read for the first `min(16, ceil(kv_len / 128))` entries of each + token; a -1 among them skips that CTA's block, entries past them are never + read. `use_pdl` only changes the launch attribute. The hand-off assumes + `num_sf_k_tiles = 8` (a 1024-wide o_proj input) and buffers with at least + as many rows as the step; `MERGE_MXFP8` sizes them at 128 rows. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + ) -> torch.Tensor: + # merge_once: this target's copy of the op (minimaxM3Sm107Tp8/sparseDecodeLc8m.cu), whose split-K merge forms + # the 16 weights once per CTA (same values, same order: bitwise the image op's output). + lc8 = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_sparse_decode_lc8 + if merge_once + else torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8 + ) + # prefetch: up to two read-only tensors this target's op pulls into L2 (a hint; the image op takes none). + # next_kv: (this layer's id buffer, the next sparse layer's, its K / V pools and block scales), target op only. + extra = () + if merge_once and (prefetch or next_kv): + extra = tuple(prefetch or ()) + (None,) * (2 - len(prefetch or ())) + if next_kv: + extra += tuple(next_kv) + lc8( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + decode_query_len, + sm_scale, + linear_softmax, + use_pdl, + prewait, + *extra, + **({"identity": True} if identity else {}), + **({"identity_step": True} if identity_step else {}), + **({"one_div": True} if one_div else {}), + ) + return out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 1024] the op is supposed to write: ``sparse_attention_reference`` in float64, + rounded once. ``linear_softmax``, ``use_pdl``, ``prewait`` and ``one_div`` do not change what is computed; + the hand-off buffers are checked separately with ``mxfp8_handoff_reference``. ``identity`` attends each token's + first min(16, blocks) blocks (its own causal length), whatever ``topk`` holds. + """ + if identity: + topk = _identity_table(seq_lens, decode_query_len, int(q.shape[0]), topk.device) + attn = sparse_attention_reference( + q.view(-1, NUM_HEADS, HEAD_DIM), + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.reshape(out.shape).to(out.dtype) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + ) -> None: + """The inputs this op takes and answers wrongly (``m3_attn_lc8.lc8_applies`` is the target's form of the + layout half). + + The kernel addresses a page's token rows at a fixed 64-byte (data) and + 8-byte (scale) pitch and loads them 16 / 8 bytes at a time, so only the + page stride is read from the views. It reads the selected block ids, + the block table and the hand-off buffers unchecked. Reads ``topk`` and + ``seq_lens`` back to the host: a test-only cost. + """ + for name, t in (("k", k), ("v", v)): + assert t.dim() == 4 and t.shape[1] == 1 and t.shape[2:] == (PAGE, HEAD_DIM // 2), ( + f"{name} must be the [pages, 1, 128, 64] byte view of one KV head" + ) + assert t.stride(3) == 1 and t.stride(2) == HEAD_DIM // 2, ( + f"{name} rows must be dense 64-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0, ( + f"{name} pages must be 16-byte aligned" + ) + for name, t in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + assert t.dim() == 4 and t.stride(3) == 1 and t.stride(2) == HEAD_DIM // SF_VEC, ( + f"{name} rows must be dense 8-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 8 == 0 and t.stride(0) % 8 == 0, ( + f"{name} pages must be 8-byte aligned" + ) + assert q.data_ptr() % 16 == 0 and q.stride(0) % 16 == 0 and q.stride(1) % 16 == 0, ( + "q rows are loaded 16 bytes at a time and must be 16-byte aligned" + ) + assert k_global_scale.numel() == 1 and v_global_scale.numel() == 1, ( + "global scales are single values" + ) + total_q = q.shape[0] + assert ( + total_q % decode_query_len == 0 and seq_lens.numel() >= total_q // decode_query_len + ), ( + "every query token must belong to a request: seq_lens is read at tok // decode_query_len unchecked" + ) + seq = seq_lens.tolist() + num_cols = block_table.shape[1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk[tok, : min(TOPK, nblocks)] + assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( + f"token {tok}: selected block ids must be -1 or below the request's block count; the block " + "table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + f"num_sf_k_tiles must be {WIDTH // 128}: the hand-off rows are {WIDTH} wide and a smaller " + "value overlaps them silently" + ) + assert ( + q8.numel() >= total_q * WIDTH + and sf is not None + and sf.numel() >= ((total_q + 127) // 128 * 128 * (WIDTH // MX_BLOCK)) + ), "the hand-off buffers must hold every row of the step: they are written unchecked" + + +def _identity_table(seq_lens, decode_query_len, total_q, device): + """[T, 16] ids (i < min(16, blocks of the token's causal length) ? i : -1), the selection identity mode reads.""" + ids = torch.full((total_q, 16), -1, dtype=torch.int32, device=device) + for tok in range(total_q): + req, qoff = divmod(tok, decode_query_len) + kv_len = max(int(seq_lens[req]) - decode_query_len + qoff + 1, 0) + n = min(16, (kv_len + PAGE - 1) // PAGE) + ids[tok, :n] = torch.arange(n, dtype=torch.int32, device=device) + return ids + + +m3_sparse_decode_lc8 = M3SparseDecodeLc8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..43166b472eca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The MSA pool's page: one sparse block is one 128-token KV page. +TOKENS_PER_BLOCK = 128 + + +class MsaKvUpdate(OpWrapper): + """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. + + ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its + own buffers at the head of every captured speculative decode step. One + single-CTA launch writes, for ``batch`` requests and ``total_q`` query + tokens (all int32): + + * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, + ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); + * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` + attends up to, through ``block_table``; + * ``n_valid_blocks[t]``, when given, the block count top-k selection is + bounded by. + + It returns None; the cells read the four buffers. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation at one request: qo = 1, every buffer written", + spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="decode without speculation at the captured batch 8", + spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", + spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), + ), + Cell( + why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", + spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), + ), + Cell( + why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", + spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), + ), + Cell( + why="positions crossing a page boundary: the slot moves to the previous page of the request", + spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), + ), + ) + + note = """ + The kernel covers any batch up to the block size's loop (it strides over + 256 threads), and any page size; only `tokens_per_block = 128`, the MSA + pool's page, is certified. + + The op does not bounds-check the request index a token reads + (`q_batch_row`) or the block-table column it derives: both come from the + metadata's own staging, which keeps them in range by construction, and an + out-of-range value reads whatever memory sits there. `is_valid` states both. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.minimax_m3_msa_kv_update( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + + def reference( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], + n_valid_blocks[:total_q])`` as the op is supposed to write them. + + The in-tree metadata's own torch re-derivation + (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook + call), statement for statement, returning new tensors rather than + writing the buffers. The buffer arguments only say which outputs exist: + a ``None`` buffer gives a ``None`` result. + """ + page = int(tokens_per_block) + kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) + cu_kv = None + if cu_kv_lens is not None: + cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) + cu_kv[1:] = torch.cumsum(kv_true, 0) + qbr = q_batch_row[:total_q].to(torch.long) + qo_dev = seq_lens[:batch] + pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(block_table.shape[1]) - 1 + ) + slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + n_valid = None + if n_valid_blocks is not None: + n_valid = torch.div( + (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" + ).to(torch.int32) + return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid + + def is_valid( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + """The inputs this op takes and answers wrongly. + + Every operand but the block table is read as a dense run from its first + element, so a strided view reads the wrong elements. The request index of + each token and the table row it selects are not bounds-checked. Reads + ``q_batch_row`` back to the host: a test-only cost. + """ + for name, t in ( + ("kv_lens", kv_lens), + ("kv_lens_staged", kv_lens_staged), + ("seq_lens", seq_lens), + ("q_batch_row", q_batch_row), + ("q_intra", q_intra), + ("msa_seq_lens", msa_seq_lens), + ("out_cache_loc", out_cache_loc), + ("cu_kv_lens", cu_kv_lens), + ("n_valid_blocks", n_valid_blocks), + ): + assert t is None or t.is_contiguous(), ( + f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" + ) + rows = q_batch_row[:total_q] + assert total_q == 0 or ( + int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) + ), ( + "q_batch_row must index a request below batch (and a block-table row): the op reads the " + "request's lengths and table row unchecked" + ) + + +msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/__init__.py new file mode 100644 index 000000000000..5da86fe24bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective-communication entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..8daa2d6b9540 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams + + +class M3AllReduce(OpWrapper): + """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the + plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: + ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; + returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the + communicator's per call. + + The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps + beyond the multicast kernel's rows. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # No cell names a world size: the result does not depend on it, and the certifying matrix runs at every world + # size one node can form (the target's two-tray group of 8 cannot be formed on one node). + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, hidden=6144)), + Cell(why="a 64-row generation step", spec=dict(m=64, hidden=6144)), + Cell(why="plain sum of eight drafter rows", spec=dict(m=8, hidden=6144, plain=True)), + ) + + note = """ + The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then + adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum + to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. + """ + + def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): + return comm(x, all_reduce_params=params) + + def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): + """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" + total = torch.stack([x.float() for x in xs]).sum(0) + if params is None: + return total.to(xs[0].dtype) + updated = (params.residual.float() + total).to(xs[0].dtype) + h = updated.float() + normed = ( + h + * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) + * params.norm_weight.float() + ) + return normed.to(xs[0].dtype), updated + + +m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..68773b5b74f2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,300 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode layer boundary, 2 ranks on SM107.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer + fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; + only the stored residual is rounded to BF16. + + The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are + summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic + and the normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + pre = total.to(torch.bfloat16).float() + residual.float() + updated = pre.to(torch.bfloat16) + u = pre.double() if drafter else updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +class M3McAllreduceNorm(OpWrapper): + """One TP group's boundary: ``updated = bf16(residual + bf16(sum_r bf16(x_r + x2_r)))``, ``normed = + RMSNorm(updated)``: bitwise the stock all-reduce + fused residual RMSNorm (+ mxfp8_quantize) of SM107 TP2. + + Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls + its own view ``ws_uc`` (re-arming the words it read) and rotates the per-CTA call parities in ``flags``; the + workspace is the ``McAllReduceWorkspace`` bucket of ``maxt`` tokens. All ranks must make the same calls in the + same stream order. ``prepushed`` 2 means a producer already pushed ``x`` (and launched this grid after its + pushes); 3 that two producers pushed ``x`` into exchange set A and ``x2`` into set B. Returns + ``[normed, updated]`` or, with ``quant``, also the MXFP8 copy of ``normed`` (E4M3 rows, swizzled UE8M0). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The tp2 decode boundaries (hidden 6144, 2 ranks, eps 1e-6): every decode step fits the 8-token bucket, epilogue + # option 7 at 1, 2, 7, 8 tokens and 15 at 3..6 (kernels/m3_mc_allreduce.AR_OPTS_BY_M; every option is + # bit-identical). + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token at the attention boundary, x only, option 7", + spec=dict(m=1, maxt=8, opts=7, x2=False, quant=False, prepushed=0), + ), + Cell( + why="MoE boundary at 4 rows (one request verifying 1 + 3): shared + routed summed in-kernel, MXFP8 copy", + spec=dict(m=4, maxt=8, opts=15, x2=True, quant=True, prepushed=0), + ), + Cell( + why="MoE boundary at 8 rows (two requests), MXFP8 copy for the next QKV", + spec=dict(m=8, maxt=8, opts=7, x2=True, quant=True, prepushed=0), + ), + Cell( + why="attention boundary whose 8 rows the o_proj GEMV already pushed (prepushed 2)", + spec=dict(m=8, maxt=8, opts=7, x2=False, quant=False, prepushed=2), + ), + # Every other decode row count (modeling's MC_ALLREDUCE), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} rows: a boundary row count between the ones above", + spec=dict( + m=m, + maxt=8, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + x2=m % 2 == 0, + quant=m % 3 == 0, + prepushed=0, + ), + ) + for m in range(1, 9) + if m not in (1, 4, 8) + ), + # The Eagle3 drafter boundaries (modeling's DRAFT_MC_AR, SpecTail): option 16 (the drafter's stock rounding) on + # 23 / 31 at every draft-pass row count 1..8 (pass 0: 4 x requests, later passes: requests), x only. + *( + Cell( + why=f"Eagle3 drafter boundary at {m} rows (opts {23 if m in (1, 2, 7, 8) else 31})", + spec=dict( + m=m, + maxt=8, + opts=23 if m in (1, 2, 7, 8) else 31, + x2=False, + quant=False, + prepushed=0, + drafter=True, + ), + ) + for m in range(1, 9) + ), + # Every other row count the o_proj pushes (modeling's ATTN_PUSH). + *( + Cell( + why=f"{m} rows prepushed by the o_proj GEMV", + spec=dict( + m=m, + maxt=8, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + x2=False, + quant=False, + prepushed=2, + ), + ) + for m in range(1, 8) + ), + ) + + note = """ + The epilogue option bits (1 early residual / weight load, 2 st.async cluster exchange, 4 early parity rotation, + 8 max carveout) are bit-identical by construction; only the four (bucket, option) pairs of the sm_103 kernel are + built. SM107 TP2 against the stock AUTO all-reduce + fused residual RMSNorm + mxfp8_quantize, rows 1..24, x alone + and x + x2: normed, updated and the MXFP8 copy (data and the scale bytes of valid rows) bitwise equal. `prepushed` + 1 (pushed, grid wait kept) is accepted by the op but no caller passes it. `pdl` changes only launch overlap and is + driven off in the cells. + + `normed` is gated to one BF16 rounding step: the kernel's RMS statistic is a warp / cluster tree of FP32 sums + and its `rsqrtf` is approximate, so a row whose scaled value lands next to a BF16 midpoint can round either + way. `updated` and the MXFP8 copy of the op's own `normed` are bit-exact. + """ + + def __call__( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + own_op: bool = False, + ) -> list[torch.Tensor]: + # The capture operands only with a capture: an op without the capture epilogue (an older image) takes the + # rest unchanged. own_op: this target's copy (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm, the same + # kernels with the capture epilogue; modeling's MC_AR_OWN_OP). + op = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_mc_allreduce_norm + if own_op + else torch.ops.trtllm.minimax_m3_sm107_tp2_mc_allreduce_norm + ) + return op( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + None, + None, + maxt, + quant, + prepushed, + opts, + *((cap_out, cap_w, cap_eps) if cap_out is not None else ()), + ) + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant, drafter) + + def is_valid( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + own_op: bool = False, + ) -> None: + h = residual.shape[-1] + # The op checks norm_w's element count, not its layout: a strided view is read as packed memory. + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + # The workspace is addressed as [set][parity][maxt][world][H / 2] words and never bounds-checked: a smaller + # buffer (another bucket's, or one built for fewer ranks) is polled and re-armed out of range. + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + # Flags are per (token, cluster CTA) of the bucket; the op checks the count, not that they belong to the + # workspace. Flags of another bucket hold another call history, and the parity they select is not the one the + # peers push into: the poll then waits on the other half forever or reads last call's rows. + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + + +m3_mc_allreduce_norm = M3McAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..8b74c9974397 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize + multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode MoE boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_mc_allreduce_norm import boundary + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class M3McFinalizeAllreduceNorm(OpWrapper): + """The MoE boundary with the routed experts' finalize folded in. + + Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, + ``scales`` fp32 or bf16), then the group's all-reduce, residual add and RMSNorm exactly as + ``m3_mc_allreduce_norm``: same workspace, flags and call-order contract. Returns ``[normed, updated]`` or, with + ``quant``, also the MXFP8 copy of ``normed``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Not taken at tp2 (modeling's MC_FINALIZE is empty: the decode step is at most 8 rows); the cells are the tp4 + # target's, on the SM107 kernel. The tp4 target takes this path from 9 tokens (top-4 of 128 experts, fp32 routing + # weights, the shared expert always present): the 16- and 48-token buckets, epilogue option 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="the smallest finalize call (9 tokens, 16-token bucket) with the next QKV's MXFP8 copy", + spec=dict(m=9, maxt=16, opts=15, topk=4, quant=True, dropped=0), + ), + Cell( + why="a full 16-token bucket with dropped expert slots", + spec=dict(m=16, maxt=16, opts=15, topk=4, quant=False, dropped=5), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=24, maxt=48, opts=15, topk=4, quant=False, dropped=0), + ), + Cell( + why="the 48-token bucket full, MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, topk=4, quant=True, dropped=3), + ), + # Every other row count of the unfinalized MoE (MC_FINALIZE): the bucket and options the target picks. + *( + Cell( + why=f"{m} tokens: a finalize row count between the ones above", + spec=dict( + m=m, + maxt=16 if m <= 16 else 48, + opts=15, + topk=4, + quant=m % 2 == 0, + dropped=m % 4, + ), + ) + for m in range(9, 49) + if m not in (9, 16, 24, 48) + ), + ) + + note = """ + BF16 routing weights and the 8-token bucket are accepted by the op (option 7 there) but the target only enters + this path at 9 tokens and up, with the built-in MoE's fp32 weights; neither is a cell. The rest is as in + `m3_mc_allreduce_norm`'s note. + """ + + def __call__( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + own_op: bool = False, + ) -> list[torch.Tensor]: + # The capture operands only with a capture: an op without the capture epilogue (an older image) takes the + # rest unchanged. own_op: this target's copy (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm, the same + # kernels with the capture epilogue; modeling's MC_AR_OWN_OP). + op = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_mc_allreduce_norm + if own_op + else torch.ops.trtllm.minimax_m3_sm107_tp2_mc_allreduce_norm + ) + return op( + fc2, + shared, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + 0, + opts, + *((cap_out, cap_w, cap_eps) if cap_out is not None else ()), + ) + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + own_op: bool = False, + ) -> None: + h = residual.shape[-1] + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + # The kernel gathers fc2[perm] without a bound: an index past the FC2 rows reads whatever follows them. + lo, hi = int(perm.min()), int(perm.max()) + assert lo >= -1 and hi < fc2.shape[0], ( + f"perm spans [{lo}, {hi}]; entries must be -1 (dropped) or a row of fc2 (< {fc2.shape[0]})" + ) + + +m3_mc_finalize_allreduce_norm = M3McFinalizeAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/__init__.py new file mode 100644 index 000000000000..e9f4d8e20b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedding ops of this target: the generic token-embedding gather.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..71fce6b253cb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="an 8-row decode verify, int64 ids", + spec=dict(rows=8, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ) + + note = """ + An exact gather: `reference` must match bit for bit. + """ + + def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) + + def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] + return rows * valid.unsqueeze(-1).to(table.dtype) + + +m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/__init__.py new file mode 100644 index 000000000000..9cc374bdf1cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMM entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..b973b175d8b4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3Bf16Linear(OpWrapper): + """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. + + The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), + o_proj (this rank's input columns), gate_up and down projections run it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, nine draft rows", spec=dict(m=9, n=1280, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=1024)), + Cell(why="drafter gate_up, 16 rows", spec=dict(m=16, n=4608, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=2304)), + ) + + note = """ + cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x, w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return (x.float() @ w.float().t()).to(x.dtype) + + +m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py new file mode 100644 index 000000000000..a0013e3ded3b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense MLP's swigluoai activation as one C++ launch with PDL and an early release (DENSE_SWIGLU_CPP).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# The dense layers' swigluoai parameters, which the kernel hard-codes (modeling_minimaxm3._build_swiglu_oai_dense_mlp). +LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 + + +class M3DenseSwiglu(OpWrapper): + """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, + rounded once to bf16: trtllm::minimax_m3_sm107_tp8_swiglu_oai (minimaxM3Sm107Tp8/swigluOai.cu), bit for bit the + Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the + grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights + during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even + half width; ``__call__`` returns None otherwise or for other activation parameters (the caller runs the Triton + op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP8 dense MLP shard: gate | up 3072 wide (d = 1536), at the decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP8 dense MLP at 1 row", spec=dict(rows=1, d=1536)), + Cell(why="the TP8 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=1536)), + Cell(why="the TP8 dense MLP at 8 rows", spec=dict(rows=8, d=1536)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the + alpha-scaled ex2.approx sigmoid, div.full), checked over every bf16 gate and every bf16 up pattern. + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float] = LIMIT, + alpha: Optional[float] = ALPHA, + beta: Optional[float] = BETA, + use_pdl: bool = True, + ) -> Optional[torch.Tensor]: + if ( + (limit, alpha, beta) != (LIMIT, ALPHA, BETA) + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 4 != 0 + or x.data_ptr() % 4 != 0 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp8_swiglu_oai(x, out, bool(use_pdl)) + return out + + def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_dense_swiglu = M3DenseSwiglu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py new file mode 100644 index 000000000000..7e0e023ee88c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 x BF16 -> BF16 low-M GEMM of the Eagle3 drafter's linears at fixed FlashInfer tactics (SM10x).""" + +from __future__ import annotations + +from typing import Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: (N, K) of the TP8 drafter shards -> ((max rows, kind, tactic), ...), rows ascending. Kind "d" is FlashInfer's direct +#: SIMT kernel with tactic (block, outputs per block, rows per block); kind "s" is its split-K tensor-core kernel with +#: tactic (mma_m, mma_n, split_k, ab_stages). A row count is served by the first entry whose max rows holds it; the +#: rows modeling's DRAFT_LOWM covers are only the ones measured (the fastest tactic at no larger error vs fp64 than +#: cuBLAS). TP8 shards (64 query / 4 KV heads of 128, KV heads replicated across rank pairs; intermediate 18432): +#: qkv 1280 x 12288, o_proj 6144 x 1024, gate_up 4608 x 6144, down 6144 x 2304. Rubin SM107 TP8 sweep 777890 (mclk +#: locked, rotating weights >= 3x L2, CUDA-graph timing), us cuBLAS -> low-M at 1 / 4 rows: qkv 6.15 -> 3.94 / 6.26 -> +#: 5.69, o 4.97 -> 2.42 / 4.48 -> 2.68, gate_up 8.09 -> 5.78 / 8.11 -> 6.99, down 5.62 -> 4.00 / 5.74 -> 3.94. The fc is +#: replicated (the TP2 shape): its TP2 tactic and row. +TACTICS: Mapping[Tuple[int, int], tuple] = { + # qkv + (1280, 12288): ((1, "d", (256, 2, 1)), (4, "d", (128, 4, 4))), + # o_proj + (6144, 1024): ((1, "s", (64, 8, 1, 10)), (4, "s", (64, 8, 2, 6))), + # gate_up + (4608, 6144): ((1, "d", (128, 2, 1)), (4, "d", (64, 4, 4))), + # down + (6144, 2304): ((1, "s", (64, 8, 2, 6)), (4, "s", (64, 8, 1, 10))), + # the Eagle3 fc over the three captured hidden states + (6144, 18432): ((1, "d", (384, 2, 1)),), +} +#: Rows each (N, K) is certified at (a subset of what its table would serve). +ROWS: Mapping[Tuple[int, int], Tuple[int, ...]] = { + (1280, 12288): (1, 4), + (6144, 1024): (1, 4), + (4608, 6144): (1, 4), + (6144, 2304): (1, 4), + (6144, 18432): (1,), +} + + +def tactic_for(n: int, k: int, m: int) -> Optional[Tuple[str, tuple]]: + """(kind, tactic) of an [m, k] x [n, k]^T call, or None when the table does not serve it.""" + for rows, kind, tactic in TACTICS.get((n, k), ()): + if m <= rows: + return kind, tactic + return None + + +class M3DraftLowM(OpWrapper): + """``out`` [M, N] BF16 = ``x`` [M, K] BF16 @ ``weight`` [N, K]^T BF16, FP32 accumulation. + + ``kind`` / ``tactic`` pick FlashInfer's kernel: "d" runs ``run_direct_dense`` (SIMT, one dot product per output in + K order), "s" runs ``run_splitk_dense`` (tcgen05 MMA, K split over ``split_k`` CTAs and summed in FP32). Both + round once to BF16. The kernel is compiled on its first call per (M, N, K, tactic, PDL); the target compiles every + covered call before warmup (``M3Eagle3Draft.prebuild_linears``). + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter qkv, 1 rows: d tactic", spec=dict(nk=(1280, 12288), m=1)), + Cell(why="drafter qkv, 4 rows: d tactic", spec=dict(nk=(1280, 12288), m=4)), + Cell(why="drafter o_proj, 1 rows: s tactic", spec=dict(nk=(6144, 1024), m=1)), + Cell(why="drafter o_proj, 4 rows: s tactic", spec=dict(nk=(6144, 1024), m=4)), + Cell(why="drafter gate_up, 1 rows: d tactic", spec=dict(nk=(4608, 6144), m=1)), + Cell(why="drafter gate_up, 4 rows: d tactic", spec=dict(nk=(4608, 6144), m=4)), + Cell(why="drafter down, 1 rows: s tactic", spec=dict(nk=(6144, 2304), m=1)), + Cell(why="drafter down, 4 rows: s tactic", spec=dict(nk=(6144, 2304), m=4)), + Cell(why="drafter fc, 1 rows: d tactic", spec=dict(nk=(6144, 18432), m=1)), + ) + + note = """ + The table is the TP8 drafter's: its shapes are the per-rank shards of the Eagle3 layer at 8-way tensor + parallelism (64 query and 4 KV heads, the KV heads replicated across rank pairs; intermediate 18432) and the + replicated fc, at the b1 decode rows (1 row in draft passes 1-2, 4 in pass 0). Certified at those shapes and + rows only; FlashInfer's kernels take other shapes, which are not measured here. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + if kind == "d": + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + DirectTactic, + run_direct_dense, + ) + + run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) + elif own_splitk: + # This target's copy of the split-K kernel (kernels/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # it also pushes the rows into a drafter boundary all-reduce (modeling's DRAFT_AR_PUSH). + from ...kernels.m3_lowm_splitk import ( + SplitKTactic, + run_splitk_dense, + run_splitk_dense_push, + ) + + if ar is not None: + run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) + else: + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + else: + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + return out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """bf16(x @ weight^T) summed in fp64 and rounded once; ``out``, ``kind``, ``tactic`` and ``use_pdl`` only + choose where and how the op writes.""" + return (x.double() @ weight.double().t()).to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> None: + """The kernels index ``x`` and ``out`` densely: a strided view is read as if it were packed.""" + assert x.dim() == 2 and x.is_contiguous(), "x must be a contiguous [M, K] matrix" + assert weight.dim() == 2 and weight.is_contiguous(), ( + "weight must be a contiguous [N, K] matrix" + ) + assert out.is_contiguous() and tuple(out.shape) == (x.shape[0], weight.shape[0]), ( + "out must be a contiguous [M, N] matrix" + ) + assert ar is None or (own_splitk and kind == "s" and len(ar) == 5), ( + "ar (the boundary all-reduce push) runs only on this target's split-K copy" + ) + + +m3_draft_lowm = M3DraftLowM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..9702cc0778a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class M3Fp32Linear(OpWrapper): + """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the + fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies + (the engine leaves torch's default).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), + Cell( + why="a generation step above the tensor-core router's 48 rows", + spec=dict(m=64, e=128, k=6144), + ), + ) + + note = """ + Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against + the true fp32 reference. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x.to(torch.float32), w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return x.float() @ w.t() + + +m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..4e9f8f65de1b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 x MXFP8 GEMV (tcgen05, block-scaled, CuTe DSL) on an activation its producer already quantized.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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 tp2, the GEMV's tuned shapes. +QKV_SPARSE = (1536, 6144) # the fused QKV + index projection (index values disabled) +QKV = (1280, 6144) +O_PROJ = (6144, 1024) +SHARED_GATE_UP = (768, 6144) +SHARED_DOWN = (6144, 384) +DENSE_GATE_UP = (3072, 6144) +DENSE_DOWN = (6144, 1536) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "shared gate_up": SHARED_GATE_UP, + "shared down": SHARED_DOWN, + "dense gate_up": DENSE_GATE_UP, + "dense down": DENSE_DOWN, +} + + +def served_cells(named: tuple[Cell, ...], x_bf16: bool, max_m: int) -> tuple[Cell, ...]: + """A cell for every (projection, rows) the GEMV serves up to ``max_m`` rows that ``named`` does not list. + + Certification is per projection: a projection's split-K, ring and MMA width differ, so each (N, K, M) the + decode dispatch can send here is its own cell. + """ + listed = {(cell.spec["nk"], cell.spec["m"]) for cell in named} + return tuple( + Cell(why=f"{name}, {m} rows", spec=dict(nk=nk, m=m)) + for name, nk in PROJECTIONS.items() + for m in range(1, max_m + 1) + if (nk, m) not in listed and _gemv.serves(nk[0], nk[1], m, x_bf16) + ) + + +def gemv_reference( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """bf16(x @ dequant(weight)^T) for fp32 activation values x [M, K], summed in fp64 and rounded once.""" + return (x.double() @ mxfp8_dequant(weight, weight_scale).double().t()).to(torch.bfloat16) + + +def gemv_is_valid(weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + n, k = weight.shape + # The kernel reads both from data_ptr() as dense buffers at the layout it was compiled for. + assert weight.dtype == torch.float8_e4m3fn and weight.is_contiguous(), ( + "weight must be contiguous e4m3 [N, K]" + ) + assert weight_scale.numel() == n * k // 32, ( + "weight_scale must be the 128x4-swizzled UE8M0 buffer [N * K / 32]" + ) + + +class M3MxGemv(OpWrapper): + """bf16 [M, N] = xq [M, K] (e4m3, 128x4-swizzled UE8M0 scales ``xsf``) x weight [N, K]^T (e4m3, swizzled + ``weight_scale``), M <= 48; None when the shape is not served (the caller falls back). + + ``trigger_early``: release the PDL dependents before the epilogue. ``cfg``: a (split, ring[, contig, nacc, a_tmem, + st_async]) tuning override; default the table's. ``use_pdl``: default the environment's. ``ar`` = (multicast + workspace address, flags address, rank, world, bucket tokens, mode): also push the rows into the next boundary + all-reduce's exchange (then ``op.last_pushed`` is True). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The projections' (N, K) at the token counts each build differs in: the MMA width (8 columns up to the TP2 decode + # step's 8 rows) is one, so one cell per projection pins the shape, and the rest are every served row count. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 8)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 8 rows: two requests verifying 1 + 3", + spec=dict(nk=QKV_SPARSE, m=8), + ), + Cell( + why="dense-attention QKV, 4 rows: one request verifying 1 + 3", spec=dict(nk=QKV, m=4) + ), + Cell(why="o_proj, one token (split 2, K 1024)", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 8 rows", spec=dict(nk=O_PROJ, m=8)), + Cell(why="shared gate_up, 4 rows (split 8)", spec=dict(nk=SHARED_GATE_UP, m=4)), + Cell(why="shared down, one token (split 1, K 384)", spec=dict(nk=SHARED_DOWN, m=1)), + Cell(why="dense gate_up, 8 rows (N 3072)", spec=dict(nk=DENSE_GATE_UP, m=8)), + Cell(why="dense down, 4 rows (K 1536)", spec=dict(nk=DENSE_DOWN, m=4)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling's MX_GEMV_ROWS). + CELLS = CELLS + served_cells(CELLS, x_bf16=False, max_m=8) + + note = """ + The kernel sums in fp32 in split-K order where `reference` sums in fp64: the bf16 outputs agree within the band the + test derives (SM107: at most as many misrounded elements as the stock FlashInfer CuTe DSL GEMM on the same + operands, rand / wide-scale / cancellation / tie / zero inputs). Not covered: the all-reduce push (`ar`), which + needs a multicast exchange; `cfg` overrides beyond the table (the target passes none). Shapes outside the table + and M > 48 return None; the target sends at most 8 rows. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv(xq, xsf, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The dequantized product; the launch options are ignored.""" + return gemv_reference(mxfp8_dequant(xq, xsf), weight, weight_scale) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + m, k = xq.shape + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous(), ( + "xq must be contiguous e4m3 [M, K]" + ) + # A linear-layout [M, K / 32] scale buffer is read as the swizzled one: wrong scales, no error. + assert xsf.numel() >= (m + 127) // 128 * 128 * (k // 32), ( + "xsf must be the 128x4-swizzled scale buffer (rows padded to 128), not the linear one" + ) + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv = M3MxGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..3b5d21235f29 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 GEMV (CuTe DSL) on a bf16 activation, quantized to MXFP8 in the kernel's prologue.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...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, + SHARED_DOWN, + SHARED_GATE_UP, + gemv_is_valid, + gemv_reference, + served_cells, +) + +__all__ = ["m3_mx_gemv_bf16"] + + +class M3MxGemvBf16(OpWrapper): + """bf16 [M, N] = mxfp8_quantize(x) x weight [N, K]^T for a contiguous, 16-byte aligned bf16 ``x`` [M, K]; the + quantization is mxfp8_quantize's (UE8M0 scale per 32 = the smallest power of two >= amax / 448). None when the + shape is not served. The other arguments as ``m3_mx_gemv``'s. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The prologue quantizes one 32-element block per thread and round: the projections at the token counts the + # target sends here (no producer-quantized activation), every TP2 decode row count. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token", spec=dict(nk=QKV_SPARSE, m=1)), + Cell(why="dense-attention QKV, 8 rows", spec=dict(nk=QKV, m=8)), + Cell(why="o_proj, 4 rows (the attention output in bf16)", spec=dict(nk=O_PROJ, m=4)), + Cell(why="shared gate_up, 8 rows", spec=dict(nk=SHARED_GATE_UP, m=8)), + Cell(why="shared down, 4 rows (K 1536)", spec=dict(nk=SHARED_DOWN, m=4)), + Cell( + why="dense gate_up, 8 rows: 24 K tiles per split-2 CTA", + spec=dict(nk=DENSE_GATE_UP, m=8), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling's MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=8) + + note = """ + As `m3_mx_gemv`. The in-kernel quantization is mxfp8_quantize's bit for bit except where the kernel's + rcp.approx(448) and an exact 1/448 put a block's amax on different sides of a power of two (not observed on the + cells). Not covered: `ar`, `cfg` overrides. Returns None for rows beyond 8 rounds of the prologue per CTA. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv_bf16(x, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The MXFP8 round trip of x times the dequantized weight; the launch options are ignored.""" + return gemv_reference(mxfp8_quant_dequant(x), weight, weight_scale) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv_bf16 = M3MxGemvBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..2bfc6ee495fb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned + +from .m3_mxfp8_linear import mxfp8_reference + + +class M3Mxfp8Gemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight + [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key + and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The decode dispatch hands it the boundary / merge MXFP8 activation when the GEMV declines rows it serves: the + # sparse-layer QKV and the o_proj at the tp8 shards. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, 4 rows", spec=dict(m=4, n=1536, k=6144)), + Cell(why="o_proj, 32 rows", spec=dict(m=32, n=6144, k=1024)), + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=384)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=384)), + ) + + note = """ + Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + """The dequantized product; ``tune`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..3f07fe2fb256 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned + +_BLOCK = 32 + + +def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" + pad_r = (rows + 127) // 128 * 128 + pad_c = (cols + 3) // 4 * 4 + s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) + return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] + + +def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, k = x.shape + s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) + factor = torch.pow(2.0, (s - 127).to(torch.float32)) + return ( + x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) + ).view(rows, k) + + +def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: + """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" + with true_fp32_matmul(): + return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=1536, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=1024, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=1536, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=768, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=1024, path="tuned")), + ) + + note = """ + Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` + quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and + multiplies dequantized operands in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + if tuned_quantize_gemm: + return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + if flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, output_dtype + ) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are + ignored.""" + m, k = x.shape + blocks = x.float().view(m, k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(-1, keepdim=True) + exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) + scale = torch.pow(2.0, exp) + q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale + with true_fp32_matmul(): + out = q.view(m, k) @ _dequant(weight, weight_scale).t() + return (out * alpha.float()).to(output_dtype) + + +m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py new file mode 100644 index 000000000000..f8deb858c843 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, +and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). + +The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values +``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 +marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a +per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: +per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 +code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved +it decodes to the BF16 weight bit for bit. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_target_head"] + +#: Verify rows per call (requests x (draft length + 1)). +MAX_ROWS = 8 +#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). +CTA_WARPS = 16 +#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. +TARGET_SHARD = 50016 +HIDDEN = 6144 +#: Codes per window: 0..14 decode, 15 marks an exception. +WINDOW = 15 +#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. +MAX_BASE = 254 - WINDOW +#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). +ROWS_PER_CHUNK = 4096 + + +def _pad_rows(weight: torch.Tensor) -> torch.Tensor: + n = weight.shape[0] + tiles = (n + 15) // 16 + if tiles * 16 == n: + return weight + return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) + + +def _fragment_view(w: torch.Tensor) -> torch.Tensor: + """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" + n, k = w.shape + # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) + return ( + w.view(n // 16, 2, 8, k // 16, 2, 4, 2) + .permute(0, 3, 2, 5, 4, 1, 6) + .reshape(n // 16, k // 16, 32, 8) + ) + + +def _exponents(weight: torch.Tensor) -> torch.Tensor: + return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF + + +def _window_base(hist: torch.Tensor) -> int: + """The base of the 15-exponent window holding the most weights (the lowest such base).""" + window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) + return int(window.argmax()) + + +def encode(weight: torch.Tensor, base: Optional[int] = None): + """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's + own window. + + Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile + << 16 | k), exc_w fp32 [E, 2] ({w, 0})). + """ + w = _pad_rows(weight) + n, k = w.shape + tiles, ks = n // 16, k // 16 + assert ks % 4 == 0, "K must be a multiple of 64" + bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF + expo = (bits >> 7) & 0xFF + if base is None: + base = _window_base(torch.bincount(expo.flatten(), minlength=256)) + code = expo - base + exc = (code < 0) | (code >= WINDOW) + code = torch.where(exc, torch.full_like(code, WINDOW), code) + sm = ((bits >> 8) & 0x80) | (bits & 0x7F) + sm = torch.where(exc, torch.zeros_like(sm), sm) + # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. + rows, cols = exc.nonzero(as_tuple=True) + order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) + rows, cols = rows[order], cols[order] + exc_w = torch.stack( + [ + w[rows, cols].float(), + torch.zeros((rows.numel(),), device=w.device), + ], + dim=1, + ) + exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) + exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) + exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( + torch.int32 + ) + # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. + smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) + cf = ( + _fragment_view(code.to(torch.uint8).view(n, k)) + .view(tiles, ks // 4, 4, 32, 8) + .to(torch.int32) + ) + # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. + cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) + for p in range(4): + cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) + cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 + # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. + seg01 = ( + smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) + ) # [.., 2, 32, 16] + seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) + packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] + return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() + + +def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): + """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" + n, _ = weight.shape + assert rows_per_chunk % 16 == 0 + hist = torch.zeros(256, dtype=torch.int64, device=weight.device) + for r in range(0, n, rows_per_chunk): + hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) + base = _window_base(hist) + packed, offs, poss, ws = [], [], [], [] + n_exc = 0 + for r in range(0, n, rows_per_chunk): + p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) + packed.append(p) + offs.append(off[:-1] + n_exc) + poss.append(pos) + ws.append(ew) + n_exc += pos.numel() + exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) + return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) + + +def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: + """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" + tiles, ks = (n + 15) // 16, k // 16 + p = packed.view(tiles, ks // 4, 3, 32, 16) + sm = ( + p[:, :, :2] + .permute(0, 1, 3, 2, 4) + .reshape(tiles, ks // 4, 32, 4, 8) + .permute(0, 1, 3, 2, 4) + .to(torch.int32) + ) + cw = ( + p[:, :, 2] + .contiguous() + .view(torch.int32) + .view(tiles, ks // 4, 32, 4) + .permute(0, 1, 3, 2) + .long() + & 0xFFFFFFFF + ) + codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( + torch.int32 + ) + bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) + bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) + frag = bits.reshape(tiles, ks, 32, 8) + # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). + w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) + return w.to(torch.int16).view(torch.bfloat16)[:n] + + +def decode( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + n: int, + k: int, + row0: int = 0, + row1: Optional[int] = None, +) -> torch.Tensor: + """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). + ``row0`` must be a multiple of 16.""" + row1 = n if row1 is None else row1 + assert row0 % 16 == 0 + t0, t1 = row0 // 16, (row1 + 15) // 16 + unit_bytes = (k // 64) * 3 * 32 * 16 + dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() + counts = (exc_off[1:] - exc_off[:-1]).long() + tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) + rows = tile_of * 16 + (exc_pos >> 16).long() + cols = (exc_pos & 0xFFFF).long() + sel = (rows >= row0) & (rows < row1) + dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) + return dec + + +def round_trip_ok( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + weight: torch.Tensor, + rows_per_chunk: int = ROWS_PER_CHUNK, +) -> bool: + """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent + window below the BF16 infinity exponent.""" + n, k = weight.shape + if not 0 <= base <= MAX_BASE or k % 64: + return False + if ( + packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 + or exc_off.numel() != (n + 15) // 16 + 1 + ): + return False + for r0 in range(0, n, rows_per_chunk): + r1 = min(n, r0 + rows_per_chunk) + dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) + if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): + return False + return True + + +class M3TargetHead(OpWrapper): + """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit + the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in + another summation order). + + ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` + (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are + the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and + # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with + # a shard whose rows are not a multiple of 16. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{m} verify row(s) of the TP4 target head shard", + spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), + ) + for m in range(1, MAX_ROWS + 1) + ), + Cell( + why="4 rows with outliers and exact zeros in many tiles (the exception lists)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), + ), + Cell( + why="2 rows over a shard of 4000 rows (250 tiles), K 1024", + spec=dict(m=2, n=4000, k=1024, exceptions=True), + ), + ) + + note = """ + The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception + adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs + from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); + integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp + order: the result is deterministic. + """ + + def __call__( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_target_head( + w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl + ) + return out + + def reference( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. + + ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" + weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + return logits.to(torch.bfloat16) + + def is_valid( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> None: + # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count + # stores a partial sum. + assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + assert out.data_ptr() != x.data_ptr(), "out must not alias x" + + +m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/__init__.py new file mode 100644 index 000000000000..77d7b84d9bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..88eb4f183b3a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 1..8 decode tokens in one CuTe DSL launch (``m3_moe``, the sm_107 TP2 build: routed +and shared expert width 1536 per rank): the router, the routed NVFP4 experts and the shared MXFP8 expert, on +caller-owned state (an ``M3MoeState`` and one ``M3MoeLayer`` per MoE layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ...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) + + +_ENGAGED: set = set() + + +class M3MoeEngine(OpWrapper): + """One MiniMax-M3 MoE layer for M (1..8) decode tokens: ``layer.state.run(layer, hidden, out, arp)``. + + ``hidden`` bf16 [M, 6144] contiguous, M = the state's token count. ``layer`` (from ``M3MoeState.layer``) carries the + TRTLLM-Gen NVFP4 expert buffers read in place, the per-expert scalars, the router weight / bias, the activation's + NVFP4 global scale and, on a ``shared`` state, the shared expert's MXFP8 weights; its state carries the compiled + kernel and the workspace every layer on it shares. Returns bf16 [M, 6144]: bf16(shared + routed) on a ``shared`` + state, else the routed partial. ``out``: write there instead of a new tensor. ``arp`` (``ar_push_desc``): also + push the rows into a multicast all-reduce exchange. + + The routed part is TRTLLM-Gen's NVFP4 MoE on the fp32 router logits: sigmoid (tanh form) + correction bias top-4 + of 128 (ties to the lower id), weights sigmoid / sum * 2.0 rounded to bf16; the activation quantized to NVFP4 with + the layer's global scale; per expert the raw FC1 sums, SwiGLU-OAI (alpha / beta / clamp), the NVFP4 requant with + fc31_scale_c, FC2, bf16(fc2_alpha x sum); the fma chain over the top-4 in order. The shared expert is the target's + MXFP8 SwiGLU-OAI MLP. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # One cell per token count: m_max is a trace-time constant, so each is its own build, and the target builds all + # eight (decode batches of 1..8 requests, spec-decode verify rows up to 8). The whys name what the build differs in. + CELLS: tuple[Cell, ...] = ( + Cell( + why="1 token: FC1 -> FC2 records replicated 32x, logits 32x, last phase without trailing FC2", + spec=dict(tokens=1), + ), + Cell( + why="2 tokens: records 32x, logits 32x, the one build without AR_PRE", + spec=dict(tokens=2), + ), + Cell(why="3 tokens: records 8x, logits 32x, 4 FC2-B builder warps", spec=dict(tokens=3)), + Cell( + why="4 tokens: one request verifying 3 drafts (the TP2 b1 point): 6-stage ring, B2IMG hand-off, " + "the 4 two-block FC2 CTAs stream no FC1", + spec=dict(tokens=4), + ), + Cell( + why="5 tokens: single records, a routing warp takes two tokens, logits 32x", + spec=dict(tokens=5), + ), + Cell( + why="6 tokens: single records, routing warps 8.. take tokens 4..5", spec=dict(tokens=6) + ), + Cell(why="7 tokens: the B2IMG global FC1 -> FC2 hand-off", spec=dict(tokens=7)), + Cell( + why="8 tokens: two requests verifying 3 drafts (the TP2 b2 point): 5-stage ring, x read from global " + "(no XSTAGE), no early group 0, the 4 two-block FC2 CTAs stream no FC1", + spec=dict(tokens=8), + ), + ) + + note = """ + The kernel's elementwise math uses the PTX approximations (rcp.approx, ex2.approx, tanh) where `reference` uses + torch's correctly rounded ones, and its GEMMs sum in tensor-core order where `reference` sums in fp64: a + quantization boundary (NVFP4 activation / intermediate code, routing near-tie) can fall on the other side, so the + two agree within the band the test derives, not bit for bit. The routed part is bit for bit the stock TRTLLM-Gen + NVFP4 MoE (fp4_quantize + fp4_block_scale_moe_runner) on the same routing; the test checks that too. + + Not covered: the all-reduce push (`arp`), which needs a multicast exchange (the decode forward's boundary all-reduce + exercises it); the non-shared state (`shared=False`), which this target never builds. Every layer on one state + must run in one stream order: the workspace is shared and tagged per call, never reset. + + No `is_valid`: the compiled call checks every buffer's dtype, shape and alignment against its signature and + raises on a mismatch (a misaligned `hidden`, an `out` that is not bf16 [M, 6144], a malformed `arp`). + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + state = layer.state + if (state.device, state.m) not in _ENGAGED: + # The engagement line (once per device and token count, on every rank: the first call, eager or captured). + _ENGAGED.add((state.device, state.m)) + logger.info( + f"[M3 sm_107 tp2] MoE engine m3_moe engaged: {state.m} tokens, {state.num_ctas} routed + " + f"{state.mod.NSH} shared CTAs, {state.mod.STAGES}-stage ring, on {state.device}" + ) + return state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Plain torch over the layer's buffers (see the class docstring); `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=layer.state.shared) + + +m3_moe_engine = M3MoeEngine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py new file mode 100644 index 000000000000..636dfcd0b614 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV of the decode MoE: ``logits[M, E] = float(x) @ W^T`` at one token.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) +#: The kernel builds each block's output scale through two ``rcp.approx.ftz.f32`` reciprocals (~2^-23 relative): +#: a value this close to an e2m1 midpoint may round to either neighbour, everything further out must match. +TIE_WINDOW = 2.0**-20 + + +def nvfp4_quantize( + x: torch.Tensor, sf_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], + the pre-rounding scaled values [M, K] that ``near_tie`` reads). + + Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. + """ + gs = sf_scale.float().reshape(()) + m, k = x.shape + blocks = x.float().view(m, k // 16, 16) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +def near_tie(scaled: torch.Tensor) -> torch.Tensor: + """Elements whose e2m1 rounding the kernel's approximate reciprocals decide (see ``TIE_WINDOW``).""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=scaled.device, dtype=torch.float32) + a = scaled.abs().unsqueeze(-1) + return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) + + +def fp32_logits(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """``float(x) @ w^T`` as a true fp32 product (TF32 off).""" + with true_fp32_matmul(): + return x.float() @ w.t() + + +class M3RouterGemv(OpWrapper): + """fp32 router logits ``float(x[M, K]) @ w[E, K]^T`` with exact fp32 FMA, one CTA per expert. + + ``x`` is BF16 with 1..16 rows, unit column stride and 16-byte rows; ``w`` is the fp32 router weight; K a + multiple of 1024. ``rows`` experts per CTA; only 1 is built. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target takes the GEMV at exactly one token (2..48 go to the tensor-core router): 128 experts, hidden 6144. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token, 128 experts x hidden 6144", + spec=dict(m=1, e=128, k=6144), + ), + ) + + note = """ + The op accepts up to 16 rows; the target sends one, so wider calls are not cells. The logits differ from a true + fp32 product only in accumulation order (per-thread FMA chains over K / 1024 strided column groups, a warp tree + and an 8-warp sum). + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_router_gemv(x, w, rows, pdl) + + def reference(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + """True fp32 product; ``rows`` and ``pdl`` do not enter the arithmetic.""" + return fp32_logits(x, w) + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> None: + # The activation rows are bulk-copied from x's base address; the op checks the row stride but not the base, + # and a misaligned bulk copy faults the context rather than raising. + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + + +m3_router_gemv = M3RouterGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py new file mode 100644 index 000000000000..0306210adaab --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV plus the routed experts' NVFP4 activation, one token.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_router_gemv import fp32_logits, nvfp4_quantize + + +class M3RouterGemvFp4(OpWrapper): + """``[logits fp32 [M, E], x_fp4 [M, K / 2] uint8, x_sf [M * K / 16] uint8]`` from one kernel. + + The logits are ``m3_router_gemv``'s; the NVFP4 activation is ``fp4_quantize(x, sf_scale, 16, False, False)`` + (16-element blocks, E4M3 scales in the linear layout, per-tensor ``sf_scale``) for the routed experts' FC1. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The fused-router decode step at one token: 128 experts, hidden 6144, the experts' input scale. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token: logits for 128 experts and the NVFP4 activation of hidden 6144", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 8.0), + ), + Cell( + why="a global scale that puts block scales near the E4M3 top, where the scale cast saturates", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 2.0), + ), + ) + + note = """ + The e2m1 codes are exact except within `TIE_WINDOW` of a rounding midpoint, where the kernel's approximate + reciprocals decide the direction. Up to 16 rows are accepted; the target sends one. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_gemv_fp4(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """True fp32 logits and the NVFP4 recipe; ``pdl`` does not enter the arithmetic.""" + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return fp32_logits(x, w), packed, sf + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + # A zero or non-finite global scale makes every block scale 0 or NaN and the data codes garbage; the op + # reads the scalar on the device and cannot refuse it. + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_gemv_fp4 = M3RouterGemvFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..f776afce11da --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TF32 tensor-core router of the decode MoE for 2..48 tokens, optionally with the NVFP4 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from .m3_router_gemv import nvfp4_quantize + + +def tf32_rna(w: torch.Tensor) -> torch.Tensor: + """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" + bits = w.contiguous().view(torch.int32) + rounded = (bits + 0x1000) & ~0x1FFF + finite = (bits & 0x7F800000) != 0x7F800000 + return torch.where(finite, rounded, bits).view(torch.float32) + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target routes 2..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. + CELLS: tuple[Cell, ...] = ( + Cell( + why="two tokens, logits only (the smallest call the target makes here)", + spec=dict(m=2, k=6144, fp4=False), + ), + Cell( + why="8 tokens with the NVFP4 activation: one 8-token MMA tile, every expert tile quantizes one row", + spec=dict(m=8, k=6144, fp4=True), + ), + Cell( + why="13 tokens: a partial second token tile", + spec=dict(m=13, k=6144, fp4=True), + ), + Cell( + why="the 48-token maximum, six token tiles, logits only", + spec=dict(m=48, k=6144, fp4=False), + ), + Cell( + why="the 48-token maximum with the NVFP4 activation", + spec=dict(m=48, k=6144, fp4=True), + ), + # Every other row count the router takes (modeling's ROUTER_TC), with and without the NVFP4 activation. + *( + Cell( + why=f"{m} rows: a router row count between the ones above", + spec=dict(m=m, k=6144, fp4=m % 2 == 1), + ) + for m in range(2, 49) + if m not in (2, 8, 13, 48) + ), + ) + + note = """ + The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both + sum the same 16 slice partials in the same order. Only the one the test device picks is driven. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> tuple[torch.Tensor, ...]: + """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" + with true_fp32_matmul(): + logits = x.float() @ tf32_rna(w).t() + if sf_scale is None: + return (logits,) + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return logits, packed, sf + + def is_valid( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + if sf_scale is not None: + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/__init__.py new file mode 100644 index 000000000000..acec8fdbb26f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalization entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..c2e14b5d8c4b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step embedding gather and the first layer's input RMSNorm in one CuTe DSL launch, into caller buffers.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm + +__all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] + +VOCAB = 200064 # MiniMax-M3's vocabulary; the decode table is the full (replicated) embedding + + +class M3EmbeddingFirstnorm(OpWrapper): + """Write ``table[ids]`` into ``raw`` (bf16 [M, 6144], a zero row for an id outside ``[0, V)``) and + ``raw * rsqrt(mean(raw^2) + eps) * weight`` into ``out``; returns ``(out, raw)`` -- the first layer's normed input + and its residual. ``ids`` int32 or int64 [M], M = 1..8; ``table`` bf16 [V, 6144] contiguous; ``weight`` bf16 + [6144]. The first call per (M, id dtype, table shape, PDL) compiles (``DecodeEmbedding`` does it at build).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target builds every (rows 1..8, id dtype) owner; the kernel is one CTA per row, so the rows differ only in + # grid size: the ends of the range and both id dtypes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, int32 ids", spec=dict(rows=1, id_dtype=torch.int32)), + Cell(why="eight requests, int32 ids", spec=dict(rows=8, id_dtype=torch.int32)), + Cell(why="one request, int64 ids", spec=dict(rows=1, id_dtype=torch.int64)), + Cell( + why="six rows (one request verifying a draft), int64 ids", + spec=dict(rows=6, id_dtype=torch.int64), + ), + # Every other row count the decode dispatch takes (modeling's EMBED_FIRSTNORM), both id dtypes. + *( + Cell( + why=f"{rows} rows: a decode row count between the ones above", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in (2, 3, 4, 5, 7) + ), + ) + + note = """ + The gathered rows are exact. The norm is flashinfer's RMSNorm arithmetic in its reduction order with a fast + rsqrt, where `reference` sums in fp64 and takes torch's rsqrt: within the band the test derives, not bit for bit. + SM107 against the stock embedding + RMSNorm module (flashinfer's kernel, 768 threads): the gathered rows are + bitwise equal; the normed rows differ in one bf16 ulp on ~5e-6 of the elements (~1.6% of rows, 1-2 elements each: + the fp32 sum-of-squares tree differs, so the rsqrt input can differ in its last bits). + An id outside [0, V) gives a zero row and a zero norm (torch's embedding would raise). No accepted input is known + to compute wrongly: the host side rejects any other dtype, shape or a strided buffer. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return m3_embed_norm(ids, table, weight, eps, raw, out, use_pdl) + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x * rstd * weight.double()).to(torch.bfloat16), rows + + +m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/__init__.py new file mode 100644 index 000000000000..d25c41dc9557 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries this target alone calls: the Eagle3 verify ops, the capture / fc_norm RMSNorms and +the step-input staging.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..be5aa46fe774 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_chunk_norm import rmsnorm_rows + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = ( + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ) + + note = """ + fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional + one-ulp move against torch's statistic. + """ + + def __call__( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) + return dst + + def reference( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" + x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h + return rmsnorm_rows(x, w, eps) + + def is_valid( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> None: + # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read + # past its elements without an error. + assert w.is_contiguous(), "w must be dense [H]" + + +m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..2f71f1b381b1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states +(``m3_spec_tail.chunk_norm``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="two requests' 8 verify rows (b2) from the capture buffer, dense [8, 3 x 6144]", + spec=dict(m=8, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ) + + note = """ + fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then + against torch's statistic; the cells' band states it. + """ + + def __call__( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) + return dst + + def reference( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" + c, h = w.shape + x = src[:, : c * h].reshape(src.shape[0], c, h) + return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) + + def is_valid( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ) -> None: + c, h = w.shape + assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + + +m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..4fe14d538a5f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat +(``m3_spec_tail.draft_input``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3DraftInput(OpWrapper): + """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = + hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of + a vocab-local embedding). Returns ``(out, residual)``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling's DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the + # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests + # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", + spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", + spec=dict(m=n, h=6144, gather=True, invalid_id=False), + ) + for n in range(1, 9) + ), + Cell( + why="third pass of 3 requests: one row each, hidden read in place", + spec=dict(m=3, h=6144, gather=False, invalid_id=False), + ), + Cell( + why="one row with an out-of-table id: the embedding half is zero", + spec=dict(m=1, h=6144, gather=False, invalid_id=True), + ), + ) + + note = """ + The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's + summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, + also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the + drafter's embedding is the whole table, so its ids are always in range. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl + ) + return out, residual + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / + ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" + valid = (ids >= 0) & (ids < table.shape[0]) + emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) + emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) + hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] + ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) + return ref, (hid.clone() if residual is not None else None) + + def is_valid( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> None: + h = table.shape[1] + # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. + assert table.is_contiguous() and hidden.is_contiguous(), ( + "table / hidden must be dense [*, H]" + ) + if rows is not None: + assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" + else: + assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" + # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + assert residual is None or ( + residual.is_contiguous() and residual.shape == (ids.shape[0], h) + ), "residual must be contiguous [M, H]" + + +m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py new file mode 100644 index 000000000000..b2487a213990 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter MLP's plain SwiGLU as one C++ launch (modeling's ``DRAFT_SILU_CPP``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3DraftSilu(OpWrapper): + """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: + trtllm::minimax_m3_sm107_tp8_silu_and_mul (minimaxM3Sm107Tp8/siluAndMul.cu), bit for bit the Triton + ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP + runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs + the Triton op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP8 drafter MLP shard: intermediate 18432 / 8 = 2304, gate | up 4608 wide, at the drafter's decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP8 drafter MLP at 1 row", spec=dict(rows=1, d=2304)), + Cell(why="the TP8 drafter MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=2304)), + Cell(why="the TP8 drafter MLP at 8 rows", spec=dict(rows=8, d=2304)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (the same PTX instruction sequence), checked over every + bf16 gate and every bf16 up pattern. + """ + + def __call__(self, x: torch.Tensor) -> Optional[torch.Tensor]: + if ( + x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 16 != 0 + or x.data_ptr() % 16 != 0 + or x.shape[0] > 65535 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp8_silu_and_mul(x, out) + return out + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul(x) + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_draft_silu = M3DraftSilu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..e63389ffc6eb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One pass of the fused greedy Eagle3 draft loop: the drafter's vocab-sharded BF16 head (streamed from its lossless +12-bit coding, ``ops.gemm.m3_target_head``, or as BF16 fragments), the cross-rank argmax and the loop's per-pass +bookkeeping in one kernel (``m3_draft_step.draft_step``).""" + +from __future__ import annotations + +import functools +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ..gemm.m3_target_head import decode as decode_coded +from .m3_shard_accept import exchange_tokens, local_argmax, make_keys + +#: Requests (head rows) per pass, and exchange rows per call half of the workspace. +MAX_TOKENS = 8 +#: 16-row weight tiles one CTA can own (the kernel's shared-memory accumulator). +MAX_TILES_PER_CTA = 32 +KERNEL_WARPS = 16 # warps per CTA (draftStep.h kWarps) + + +@functools.lru_cache(maxsize=64) +def max_tile_pieces(n_tiles: int, num_ctas: int, chunks_per_tile: int) -> int: + """The most warps that stream one weight tile under the kernel's whole-tile split: CTA c owns tiles + [c * n / grid, (c + 1) * n / grid), and its (tile, chunk) range is cut into KERNEL_WARPS equal contiguous + segments. The shared-memory atomics add a tile's pieces in arrival order: deterministic for at most 2 pieces (two + addends from 0 commute), not for more (fixed_order then sums them in warp order).""" + worst = 1 + for c in range(num_ctas): + lo, hi = c * n_tiles // num_ctas, (c + 1) * n_tiles // num_ctas + span = (hi - lo) * chunks_per_tile + bounds = [w * span // KERNEL_WARPS for w in range(KERNEL_WARPS + 1)] + for t in range(hi - lo): + a, b = t * chunks_per_tile, (t + 1) * chunks_per_tile + worst = max( + worst, sum(1 for w in range(KERNEL_WARPS) if bounds[w] < b and bounds[w + 1] > a) + ) + return worst + + +#: Draft length of the fused loop, and the drafter's TP8 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 25008 +HIDDEN = 6144 + + +def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the + 12-bit coded stream with its exception lists decoded.""" + if exc_off is None: + return from_fragments(wf) + n = (exc_off.numel() - 1) * 16 + return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + + +def shard_logits( + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() + + +class M3DraftStep(OpWrapper): + """One fused greedy draft pass over ``rows <= 8`` requests (one head row each). + + Head input ``x[x_rows[b]]`` (pass 0: the last accepted verify row of each request) or ``x[b]``. The head is this + rank's BF16 lm_head shard in fragment order (``m3_draft_step.bf16_fragments``, global ids ``vocab_offset ..``), + or with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base`` its lossless 12-bit coding (``wf`` = the uint8 stream + of ``m3_target_head.encode_chunked``, which must round-trip to the BF16 shard); each logit is the fp32 dot product + rounded to BF16 and the row's token is the global argmax over all ranks' + shards (ties, -0 / +0 included, to the lowest global id), exchanged through the multicast slots + ``[2][8][world]``. The last CTA then writes ``draft_out[b, step]`` and ``ids_next[b]`` (the token), ``pos_next[b] + = pos_src[x_rows[b] or b] + 1``, and advances the draft KV lengths as the in-tree draft loop does between passes + (pass 0 with ``num_acc``: ``kv_lens -= draft_len - num_acc`` and ``seq_lens = 1``; later passes: ``kv_lens += + 1``). ``out_keys`` (optional) receives each row's winning key. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling's DRAFT_STEP): draft_len 3 at the TP8 GEN worker's 1..2 requests, + # the drafter's full-vocabulary BF16 head (25008 rows per rank of 8, hidden 6144; 1563 tiles, 8 per CTA on 212 + # SMs). Pass 0 reads the verify rows' hidden (4 per request) through x_rows and rolls the KV back; the later passes + # read one row per request. BF16 fragments only: on sm_107 the 12-bit coded head is slower than the BF16 stream + # (modeling's DRAFT_HEAD_CODED is False). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"pass 0 at {n} request(s): head rows gathered from {4 * n} verify rows, KV rollback, seq_lens = 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=0, + rank=7 * (n - 1), + world=8, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ) + for n in (1, 2) + ), + Cell( + why="pass 1 at two requests: one row each, kv_lens + 1, rank 6's ids", + spec=dict( + nreq=2, + draft_len=DRAFT_LEN, + step=1, + rank=6, + world=8, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ), + Cell( + why="last pass (step 2) at one request, rank 0", + spec=dict( + nreq=1, + draft_len=DRAFT_LEN, + step=2, + rank=0, + world=8, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ), + Cell( + why="pass 2 at 2 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=2, + draft_len=DRAFT_LEN, + step=2, + rank=3, + world=8, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + coded=False, + ), + ), + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit; + the coded head's weight is the BF16 shard bit for bit; its exceptions decode to +0 in the MMAs and their exact + products are added at the tile's flush) + and the argmax is the in-tree one over BF16 logits, so a token differs from the in-tree path's only where the fp32 + sum, accumulated in another order, rounds to the other side of a BF16 boundary and that decides the row's maximum. + The cells' rows are decidable and the peers' keys sit 0.25 away (or tie exactly); the edge tests (forced BF16 + ties, -0 / +0, top-2 one fp32 step apart that round to one BF16 value, adjacent BF16 values) use integer data + whose fp32 sums are exact in any order and hold the kernel to the in-tree path bit for bit. The cross-rank + exchange is certified on one device with the peers' keys written into the slots (see m3_shard_accept). + """ + + def __call__( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> None: + # fold: (the 14 fold-tail tensors in buildParams' order, rows, tokens per step) -- the step's forward-tail store + # and next-step staging, run by the pass's last CTA after its tokens (modeling's DRAFT_STEP_FOLD_TAIL). + tensors, fold_rows, fold_t = ( + (list(fold[0]), int(fold[1]), int(fold[2])) if fold is not None else (None, 0, 0) + ) + torch.ops.trtllm.minimax_m3_sm107tp8_draft_step( + x, + x_rows, + rows, + wf, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + num_ctas, + pdl, + out_keys, + exc_off, + exc_pos, + exc_w, + exp_base, + tile_lo, + tile_hi, + total_ctas, + None, + None, + tensors, + fold_rows, + fold_t, + fixed_order, + ) + + def reference( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> dict[str, torch.Tensor]: + """Every buffer the pass writes, computed on copies: "logits" (this rank's BF16 logits as fp32), "local_keys" + (this rank's key per row), "keys" (the winning key), "tokens", "draft_out", "ids_next", "pos_next", "kv_lens" + and "seq_lens". + + Reads ``ws_uc``, ``counters`` and ``kv_lens`` before the call. ``keys``, ``ws_mc_ptr``, ``num_ctas``, ``pdl``, + ``out_keys`` and ``fixed_order`` do not change the result; ``fold`` adds the fold tail's writes (m3_fold_tail), + not these. + """ + h = x[x_rows.long()] if x_rows is not None else x[:rows] + logits = shard_logits(h, wf, exc_off, exc_pos, exc_w, exp_base) + best, idx = local_argmax(logits, vocab_offset) + local_keys = make_keys(best, idx) + tokens = exchange_tokens(local_keys, ws_uc, counters, rank, world, MAX_TOKENS) + half = int(counters[1].item()) & 1 + slots = ws_uc.view(torch.int64)[ + half * MAX_TOKENS * world : (half * MAX_TOKENS + rows) * world + ] + slots = slots.view(rows, world).clone() + slots[:, rank] = local_keys + sign = torch.iinfo(torch.int64).min + win = (slots ^ sign).max(dim=1).values ^ sign + d_out = draft_out.clone() + d_out[:rows, step] = tokens + src = x_rows.long() if x_rows is not None else torch.arange(rows, device=x.device) + lens = kv_lens.clone() + sl = seq_lens.clone() if seq_lens is not None else None + if step == 0 and num_acc is not None: + lens[:rows] -= draft_len - num_acc[:rows] + sl[:rows] = 1 + elif step > 0: + lens[:rows] += 1 + out = { + "logits": logits, + "local_keys": local_keys, + "keys": win, + "tokens": tokens, + "draft_out": d_out, + "ids_next": tokens.clone(), + "pos_next": pos_src[src] + 1, + "kv_lens": lens, + } + if sl is not None: + out["seq_lens"] = sl + return out + + def is_valid( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> None: + assert not bool(keys[:rows].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + # The kernel reads the head row x + x_rows[b] * K and the position pos_src[x_rows[b]] with no bound. + if x_rows is not None: + assert bool(((x_rows[:rows] >= 0) & (x_rows[:rows] < x.shape[0])).all()), ( + "x_rows must index x" + ) + assert bool((x_rows[:rows] < pos_src.shape[0]).all()), ( + "pos_src must cover the head rows" + ) + else: + assert pos_src.shape[0] >= rows, "pos_src must cover the head rows" + # Each CTA accumulates at most MAX_TILES_PER_CTA tiles in shared memory. + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + assert -(-tiles // num_ctas) <= MAX_TILES_PER_CTA, "too many weight tiles per CTA" + # D-FIX: the atomics' arrival-order sum is deterministic only while no tile takes more than 2 warp pieces. + assert ( + fixed_order + or max_tile_pieces(tiles, num_ctas, x.shape[1] // (16 if exc_off is None else 64)) <= 2 + ), ( + "a tile takes more than 2 warp pieces on this grid: the atomics' sum is order-dependent (fixed_order)" + ) + if exc_off is not None: + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), ( + "exc_off must end at the exception count" + ) + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + # A NaN logit maps to a key above every finite one and wins the row silently. + h = x[x_rows.long()] if x_rows is not None else x[:rows] + assert not bool(torch.isnan(h).any()), "NaN head inputs win the argmax" + if fold is not None: + assert len(fold) == 3 and len(fold[0]) == 14, ( + "fold is (the 14 fold-tail tensors, rows, tokens per step)" + ) + if num_acc is not None: + na = num_acc[:rows] + assert bool(((na >= 1) & (na <= draft_len + 1)).all()), ( + "num_acc counts must lie in [1, draft_len + 1]" + ) + + +m3_draft_step = M3DraftStep() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..2001bf183b54 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ...kernels.m3_force_vec import force_vec_kernel, force_vec_req_kernel + + +class M3ForceVec(OpWrapper): + """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` + (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 + ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + force_vec_kernel[(1,)]( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + """The counts (a new tensor; ``counter`` and ``out`` are not written).""" + c = int(counter.item()) + 1 + slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) + idx = (c * counter_stride + slots * slot_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec = M3ForceVec() + + +class M3ForceVecReq(OpWrapper): + """For each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(keys[b] * key_stride) & pool_mask] < frac), max_total)`` into ``out`` (int32): + one draw per request at its own key (its KV length at the verify), no shared counter. ``block`` is the launch + width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="two requests, fraction 0.3", spec=dict(batch=2, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + force_vec_req_kernel[(1,)]( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + """The counts (a new tensor; ``out`` is not written).""" + k = keys[:batch_size].to(torch.int64) + r = pool[(k * key_stride) & pool_mask] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec_req = M3ForceVecReq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..d86ee7e23625 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampled (temperature + top-p) Eagle3 verification over the TP-sharded vocabulary in one kernel +(``m3_sampled_accept.sampled_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_shard_accept import strict_accept + +#: Verify rows per call. +MAX_ROWS = 48 +#: The fused sampling kernel's 1 / T guard. +TEMP_EPSILON = 1e-6 + + +def nucleus_probs( + logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +class M3SampledAccept(OpWrapper): + """Sampled verification of ``rows = nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard. + + Each row draws one token from its top-p nucleus of softmax(logits / T) over the full (all-rank) vocabulary by + distributed rejection sampling (Philox keyed by the row's seed / offset, the row, the attempt and the CTA), then + the strict acceptance runs against ``draft`` as in ``m3_shard_accept``: ``acc_tokens``, ``num_acc``, + ``gather_ids``, optional ``next_new[:, 0]``. ``tokens`` receives the sampled token per row; ``extra_attempts`` + accumulates exchange rounds beyond the second. The exchange goes through the multicast workspace (two call + halves of ``minimax_m3_sampled_accept_ws_words(world)`` words) with the (finished rows, call) ``counters``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3 (T = 4 rows per request) at the TP8 GEN worker's batch: 1 / 2 requests (4 / 8 rows, 8 CTAs per + # row), and the 48-row maximum (4 CTAs per row, 12.5 KB of staged logits per CTA); the 25008-column target shard + # (vocab 200064 / 8). Single rank (see note). + CELLS: tuple[Cell, ...] = ( + Cell( + why="b1 (4 rows), top_p 0.8 over a 12-token head: the draws follow the nucleus distribution", + spec=dict( + nreq=1, vs=25008, vocab_offset=0, check="distribution", top_p=0.8, temperature=0.7 + ), + ), + Cell( + why="b2 (8 rows): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=2, + vs=25008, + vocab_offset=7 * 25008, + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ), + Cell( + why="b2 (8 rows) with top_p = 1 (no cut: the first draw stands), next_new and forced counts", + spec=dict( + nreq=2, + vs=25008, + vocab_offset=7 * 25008, + check="distribution", + top_p=1.0, + temperature=1.0, + force="vec", + ), + ), + Cell( + why="AgentX sampling (T 1.0, top_p 0.95) at b1 with forced counts: the draws follow the nucleus", + spec=dict( + nreq=1, + vs=25008, + vocab_offset=0, + check="distribution", + top_p=0.95, + temperature=1.0, + force="total", + ), + ), + Cell( + why="48 rows (4 CTAs per row): the second launch geometry, nucleus distribution at top_p 0.5", + spec=dict( + nreq=12, + vs=25008, + vocab_offset=0, + check="distribution", + top_p=0.5, + temperature=1.3, + force="total", + ), + ), + ) + + note = """ + What is sampled is gated statistically, the rest exactly. Every drawn token must lie in the reference nucleus; + the empirical distribution over many calls (distinct Philox offsets) must sit within a stated total-variation + band of the reference's; the acceptance outputs must equal the strict acceptance of the drawn tokens exactly. + The Philox stream itself is not reproduced in torch: the kernel's draw uses the hardware exp / log + approximations (__expf, __logf) and a CTA-order inverse CDF, which torch does not reproduce bit for bit. + + Certified at world 1 only. The decision is replicated from every rank's per-round records, which another rank + computes from its own shard on the fly, so peers cannot be stood in for on one device; the cross-rank path + needs the multi-device job and is not covered here. Rows with top-k or min-p must not reach this op (it applies + neither); is_valid states it. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept( + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + rows = logits.shape[0] + # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the + # wrong rows. + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + # Per-row parameters are read at the row index, unchecked. + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), ( + "per-row parameters must cover every row" + ) + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + # The finished-row counter elects the CTA that runs the acceptance. + assert int(counters[0].item()) == 0, ( + "counters[0] (finished rows) must be zero between calls" + ) + + +m3_sampled_accept = M3SampledAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..3c3974cca0d8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept_ws_words.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Size of the sampled-verify exchange workspace (``m3_sampled_accept.SampledAcceptWorkspace.create``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Exchange rounds per row, verify rows per call, CTAs per row (at most), uint64 words per record. +MAX_ATTEMPTS = 16 +MAX_ROWS = 48 +MAX_SPLITS = 8 +RECORD_WORDS = 4 +#: SA-1R (one_round_k 9 / 82): 9 listed bf16 values ride each record, 3 per extra word. +LIST_WORDS = {0: 0, 9: 3, 82: 3} + + +class M3SampledAcceptWsWords(OpWrapper): + """uint64 words of the multicast workspace ``minimax_m3_sm107tp8_sampled_accept`` exchanges through for + ``world`` ranks: two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="the tp8 target's workspace", spec=dict(world=8)), + Cell(why="a single rank (local workspace, no multicast)", spec=dict(world=1)), + Cell( + why="the tp8 target's workspace with SA-1R pooled (82)", + spec=dict(world=8, one_round_k=82), + ), + Cell( + why="the tp8 target's workspace with SA-1R top-9 (9)", spec=dict(world=8, one_round_k=9) + ), + ) + + note = """ + A host computation (no device work); it is an op so the workspace size and the kernel's indexing come from one + definition. + """ + + def __call__(self, world: int, one_round_k: int = 0) -> int: + return torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words(world, one_round_k) + + def reference(self, world: int, one_round_k: int = 0) -> int: + """Two halves of attempts x rows x (world x splits) records x record words.""" + return ( + 2 + * MAX_ATTEMPTS + * MAX_ROWS + * world + * MAX_SPLITS + * (RECORD_WORDS + LIST_WORDS[one_round_k]) + ) + + +m3_sampled_accept_ws_words = M3SampledAcceptWsWords() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..6df70b3f467b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Greedy Eagle3 verification from this rank's logits shard: per-row argmax, the cross-rank (value, index) exchange +and the strict acceptance in one kernel (``m3_shard_accept.shard_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Verify rows per call, and exchange rows per call half of its workspace. +SHARD_ROWS = 48 + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out + + +class M3ShardAccept(OpWrapper): + """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` + (global ids ``vocab_offset ..``). + + Each row's token is its global argmax over all ranks' shards: keys ``ordered(value) << 32 | ~index`` (ties take + the lowest index, as ``torch.argmax``) are max-reduced over the shard, pushed into every rank's exchange slot + through the multicast slots ``[2][48][world]`` (Lamport: 0 = empty) and max-reduced over the ranks by the last + CTA, which then runs the strict acceptance: ``acc_tokens [nreq, D + 1]``, ``num_acc [nreq]`` (``force_vec`` / + ``force_total`` override the counts), ``gather_ids [nreq]`` (the first draft pass's row per request) and, when + given, ``next_new[:, 0]`` (the last accepted token). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every greedy verify step of the target (modeling's SHARD_ACCEPT): draft_len 3 at the TP8 GEN worker's 1..2 + # requests (4..8 rows), plus the 48-row maximum; the target lm_head's BF16 shard logits (25008 columns per rank + # at TP8, 8 exchange slots per row). The logits come from the head GEMM, so the op only reduces. + CELLS: tuple[Cell, ...] = ( + Cell( + why="b1: 4 verify rows (1 request x T = 4), rank 0 of 8, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=25008, world=8, rank=0, force="none", strided=False), + ), + Cell( + why="b2: 8 verify rows, rank 7, global ids offset by 7 x 25008", + spec=dict(nreq=2, draft_len=3, vs=25008, world=8, rank=7, force="none", strided=False), + ), + Cell( + why="8 rows (2 requests) read from a row-strided view of a wider logits buffer, forced counts", + spec=dict(nreq=2, draft_len=3, vs=25008, world=8, rank=3, force="vec", strided=True), + ), + Cell( + why="the 48-row maximum (12 requests at draft 3), force_total, rows with exact cross-rank ties", + spec=dict( + nreq=12, + draft_len=3, + vs=25008, + world=8, + rank=5, + force="total", + strided=False, + ties=True, + ), + ), + ) + + note = """ + The op reduces BF16 logits exactly, so the reference is bit-exact, including cross-rank ties (lowest global + index wins). The cross-rank exchange is certified on one device: the other ranks' keys are written into this + rank's slots before the call, which is what their multicast pushes would have put there; the multicast mapping + itself (ws_mc_ptr = this rank's slots, as at world 1) is not exercised across devices here. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + num_ctas, + force_vec, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Lowest-index maximum of each row's BF16 logits, the largest key among it and the peers' keys in + ``ws_uc``, then the strict acceptance. Reads ``ws_uc`` / ``counters`` before the call; the scratch and + output buffers and ``num_ctas`` do not change the result.""" + best, idx = local_argmax(logits.float(), vocab_offset) + tokens = exchange_tokens(make_keys(best, idx), ws_uc, counters, rank, world, SHARD_ROWS) + out = strict_accept(tokens, draft, batch_idx, force_total, next_new, force_vec) + out["tokens"] = tokens + return out + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[: logits.shape[0]].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + assert ws_uc.numel() >= 2 * SHARD_ROWS * world, "ws_uc must hold [2][48][world] slots" + # A NaN logit maps to a key above every finite one and wins the row silently. + assert not bool(torch.isnan(logits).any()), "NaN logits win the argmax" + nreq = draft.shape[0] + assert bool(((batch_idx[:nreq] >= 0) & (batch_idx[:nreq] < nreq)).all()), ( + "batch_idx must index the nreq requests" + ) + if force_vec is not None: + fv = force_vec[:nreq] + assert bool(((fv >= 1) & (fv <= draft.shape[1] + 1)).all()), ( + "force_vec counts must lie in [1, D + 1]" + ) + + +m3_shard_accept = M3ShardAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..ff07c54de1f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> +``m3_spec_tail.slot_scatter``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: + """src rows zero-padded (or cut) to ``width`` columns.""" + out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) + w = min(width, src.shape[1]) + out[:, :w] = src[:, :w] + return out + + +class M3SlotScatter(OpWrapper): + """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, + ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The + destinations are the store's transposed [slots, width] views.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of + # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) + # is zero-padded. + CELLS: tuple[Cell, ...] = ( + Cell( + why="8 rows into shuffled slots, widths equal to the store's", + spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), + ), + Cell( + why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", + spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), + ), + ) + + note = """ + Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. + """ + + def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): + torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) + + def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): + """The four destinations after the call, on copies.""" + s = slots.long() + out = {} + for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): + r = dst.clone() + r[s] = padded_rows(src[: s.numel()], dst.shape[1]) + out[name] = r + lr = dl.clone() + lr[s] = sl[: s.numel()] + out["dl"] = lr + return out + + def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): + assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( + "slots must index the store" + ) + # Two rows into one slot race (which one lands is unordered). + assert slots.unique().numel() == slots.numel(), "slots must be distinct" + assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + + +m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py new file mode 100644 index 000000000000..56c5f54b1161 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's speculative step correction as one launch (``RuntimeHooks.spec_offsets_op``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger + +# The engagement line is logged once per process (module state: __call__ reads no instance state). +_ENGAGED: list = [] + + +class M3SpecOffsets(OpWrapper): + """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp8_spec_offsets (minimaxM3Sm107Tp8/specOffsets.cu), the + same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other + layouts (the caller runs the ATen adds).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="1 request, 4 tokens (draft length 3)", spec=dict(tokens=4, seqs=1)), + Cell(why="8 requests, 32 tokens", spec=dict(tokens=32, seqs=8)), + Cell(why="wrap-around values", spec=dict(tokens=8, seqs=2, wrap=True)), + ) + + note = "Integer adds: the reference (the two in-place ATen adds) is exact." + + def __call__(self, pos, pos_off, kv, kv_off) -> bool: + ts = (pos, pos_off, kv, kv_off) + if ( + any( + t.dtype != torch.int32 or t.dim() != 1 or t.stride(0) != 1 or not t.is_cuda + for t in ts + ) + or pos.numel() != pos_off.numel() + or kv.numel() != kv_off.numel() + ): + return False + torch.ops.trtllm.minimax_m3_sm107_tp8_spec_offsets(pos, pos_off, kv, kv_off) + if not _ENGAGED: + _ENGAGED.append(True) + logger.info("[M3 host] spec step correction: one launch taken (spec_offsets_op)") + return True + + def reference(self, pos, pos_off, kv, kv_off) -> None: + pos += pos_off + kv += kv_off + + def is_valid(self, pos, pos_off, kv, kv_off) -> None: + assert pos.numel() == pos_off.numel() and kv.numel() == kv_off.numel(), ( + "equal lengths per pair" + ) + + +m3_spec_offsets = M3SpecOffsets() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..a0da8b337574 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch +(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS: tuple[Cell, ...] = ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + Cell( + why="an unchanged batch of 3 previous requests (c3) in shuffled slots", + spec=dict(p=3, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="an unchanged batch of 4 previous requests (c4) in shuffled slots", + spec=dict(p=4, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + ) + + note = """ + Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid + states them. + """ + + def __call__( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ) + + def reference( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" + nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) + s = slots[:P].long() + ids_r, draft_r = ids.clone(), draft.clone() + ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) + draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) + pos_r = torch.zeros_like(pos_off) + pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] + kv_r = torch.zeros_like(kv_off) + kv_r[b0 : b0 + P] = lens[s] - T + return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} + + def is_valid( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or + # writes another buffer. + s = slots[:P] + assert bool( + ( + (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) + ).all() + ), "slots must index the store" + assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( + "ppi must index lens" + ) + assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" + assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" + assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( + "offset buffers too short" + ) + + +m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py new file mode 100644 index 000000000000..a6c626ff2d2f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode graph's step-input prologue in one launch (this target's ``RuntimeHooks.pack_prologue_op``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepPrologue(OpWrapper): + """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch + (trtllm::minimax_m3_sm107_tp8_step_prologue, minimaxM3Sm107Tp8/stepPrologue.cu). + + ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one + ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the + window, a device copy's a device address). The launch writes up to ``max_desc`` copies of each kind, every host + copy before any device copy, then zeroes the header; with ``epoch`` it increments ``epoch[0]`` once. The default + prologue does the same with two parallel-scatter launches (``step_inputs._graph_scatter_kernel``) and a header + fill.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP8 decode step's packs: a few dozen host copies of 4 to a few hundred bytes (positions, KV lengths, slots, + # the spec sampler's uploads), a handful of device copies (the graph runner's static inputs), some reading what a + # host copy of the same window wrote. + CELLS: tuple[Cell, ...] = ( + Cell(why="an empty window (a replay with no new upload)", spec=dict(h2d=0, d2d=0)), + Cell(why="host copies only, odd sizes and alignments", spec=dict(h2d=40, d2d=0)), + Cell( + why="host copies feeding device copies (a device copy reads a host copy's destination)", + spec=dict(h2d=24, d2d=8, chained=True), + ), + Cell( + why="every descriptor slot (128 + 128 of the pack's 256)", spec=dict(h2d=128, d2d=128) + ), + ) + + note = """ + Byte copies: the reference (the default prologue's Triton launches and fill on the same window) is exact. + Destinations are raw addresses; is_valid cannot check them. + """ + + def __call__(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + torch.ops.trtllm.minimax_m3_sm107_tp8_step_prologue(window, epoch, max_desc) + + def reference(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + """The default prologue on the same arguments: it writes the same destinations.""" + from tensorrt_llm._torch.pyexecutor.step_inputs import _graph_scatter_kernel + + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + bump = epoch is not None and not device_src + _graph_scatter_kernel[(max_desc,)]( + window, + stage_i64, + epoch if bump else stage_i64, + DEVICE_SRC=device_src, + BLOCK=256, + BUMP_EPOCH=bump, + num_warps=1, + ) + window[:16].zero_() + + def is_valid(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + assert window.dtype == torch.uint8 and window.is_contiguous(), "window: contiguous uint8" + assert window.numel() >= 16 + 24 * max_desc, "window: header and max_desc descriptors" + counts = window[:16].view(torch.int64).tolist() + assert counts[0] + counts[1] <= max_desc, "at most max_desc copies in all" + + +m3_step_prologue = M3StepPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/weights.py new file mode 100644 index 000000000000..871d42635e24 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/weights.py @@ -0,0 +1,884 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_107 / tp8. + +Two tables, one per checkpoint the target reads: + +* ``MODEL_WEIGHTS`` (``MiniMaxM3Tp8Weights``): the vision-language checkpoint's language model, + ``language_model.*``. Every rank is handed the whole checkpoint, so every tensor-parallel split is this + table's: the Q heads, the K/V/index-Q heads (4 each over 8 ranks: rank r holds head r // 2, as the built-in + loader duplicates them), the dense MLPs and the shared expert's + intermediate, the o_proj / down-projection input columns, every routed expert's intermediate (all 128 + experts live on every rank), and the vocabulary of the shell's ``lm_head``. The index-K head is replicated. +* ``DRAFT_WEIGHTS`` (``MiniMaxM3Tp8DraftWeights``): the one-layer BF16 Eagle3 drafter's own checkpoint + (``speculative_model``), loaded through the shell's ``load_draft_weights`` when speculation is on. + +Storage is HF ``[out, in]`` row-major. Five families need a transform, each a layout a kernel demands: + +* **MXFP8 block scales** (every attention projection, the dense MLPs and the shared expert): the checkpoint + stores UE8M0 bytes ``[N, K / 32]``; the GEMMs read the 128x4-swizzled order + (``trtllm::block_scale_interleave``). A fused projection concatenates its shards' scales first and swizzles + once. +* **The five-way sparse-layer projection** ``[Q | K | V | index-Q | index-K]`` and the three-way dense one + ``[Q | K | V]``: the per-rank shards concatenated along N. +* **Gated MLPs** (dense and shared expert): ``[gate ; up]`` rows. +* **NVFP4 routed experts** (layers 3..59): ``[up (w3) ; gate (w1)]`` rows, then the trtllm-gen row shuffle of + the weight and of its E4M3 scales (gated-activation reorder + 128-row epilogue tile shuffle, + ``trtllmgen_maybe_get_cached_*_permute_indices``) and the 128x4 swizzle of the scales. +* **Gemma boundary norms**: ``input_layernorm``, ``post_attention_layernorm`` and ``model.norm`` are stored as + Gemma norms, ``(1 + w) * x``; the fused all-reduce + RMSNorm kernels apply ``w * x``, so ``1`` is added + (in BF16) at load. The per-head Q/K/index norms stay Gemma: their kernels apply the offset themselves. + +The routed experts' per-tensor scalars (input scales and second-level weight scales, every expert) load raw; +``modeling`` derives the kernels' global scales, alphas and activation constants from them once the weights +are real. + +Loading contract: ``load(model, weights)`` consumes the engine-provided checkpoint dict, fills every declared +parameter exactly once and asserts coverage in both directions: every declared parameter written, and every +checkpoint key either consumed or in ``expected_unconsumed`` (the vision tower and its projector, which this +text-only target does not build, and the language model's multi-token-prediction head). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) + +#: ``layers=ALL_LAYERS`` is the common case: a weight every decoder layer carries. +ALL_LAYERS = object() + +_LM = "language_model." +#: Checkpoint key families this target does not build: the vision tower and its two projectors (text only), and the +#: language model's multi-token-prediction head (speculation, when on, is the separate Eagle3 drafter). +_UNBUILT_PREFIXES = ( + "vision_tower.", + "multi_modal_projector.", + "patch_merge_mlp.", + _LM + "model.mtp.", +) + +MX_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 elements along K +NV_BLOCK = 16 # NVFP4: one E4M3 scale per 16 elements along K +EPILOGUE_TILE_M = ( + 128 # the trtllm-gen MoE runner's row tile: fixes its weight and scale row shuffles +) + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + ``shape`` takes the dimension bundle and returns the declared shape. ``src`` is a callable + ``(d, i) -> [(key or key tuple, index into the param | None), ...]`` with keys relative to the checkpoint + root; a key tuple hands the transform those tensors together. ``layers`` is ``ALL_LAYERS``, ``None`` (not + per-layer) or ``d -> layer indices``. ``transform(src, d)`` runs on the destination's device before the copy. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + +# ------------------------------------------------------------------------------------------------ shards + + +def _col( + t: torch.Tensor, d: SimpleNamespace, tp: int | None = None, rank: int | None = None +) -> torch.Tensor: + """This rank's rows (column-parallel output split), even split, sliced on the host and moved to the + destination device (only this rank's bytes are read).""" + tp = d.tp if tp is None else tp + rank = d.rank if rank is None else rank + n = t.shape[0] // tp + return _dev(t[rank * n : (rank + 1) * n], d) + + +def _heads(t: torch.Tensor, d: SimpleNamespace, head_dim: int) -> torch.Tensor: + """This rank's rows of a projection split by heads of ``head_dim`` rows: an even split when there are at least as + many heads as ranks, else the head of rank // (ranks per head), duplicated on those ranks.""" + n = t.shape[0] // head_dim + if n >= d.tp: + return _col(t, d) + h = d.rank // (d.tp // n) + return _dev(t[h * head_dim : (h + 1) * head_dim], d) + + +def _row(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """This rank's input columns (row-parallel input split), even split, on the destination device.""" + n = t.shape[1] // d.tp + return _dev(t[:, d.rank * n : (d.rank + 1) * n], d) + + +def _dev(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """``t`` on the destination device (``d.dev``, set per fill).""" + return t.to(d.dev, non_blocking=True) + + +def _swizzle(scale: torch.Tensor) -> torch.Tensor: + """128x4 swizzle of a row-major ``[N, K / 32]`` UE8M0 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(scale.contiguous()) + + +# --------------------------------------------------------------------------------------------- transforms + + +def _gemma_fold(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A Gemma boundary norm as the plain norm the fused all-reduce epilogue applies: ``w + 1`` in BF16.""" + return t + 1.0 + + +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) + + +def _qkv_sparse_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[Q | K | V | index-Q | index-K]``: Q by attention heads, K/V/index-Q by their heads, index-K replicated.""" + q, k, v, iq, ik = t + hd, xd = d.head_dim, d.index_dim + return torch.cat( + [_col(q, d), _heads(k, d, hd), _heads(v, d, hd), _heads(iq, d, xd), _dev(ik, d)], dim=0 + ).contiguous() + + +def _qkv_sparse_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v, iq, ik = t + hd, xd = d.head_dim, d.index_dim + return _swizzle( + torch.cat( + [_col(q, d), _heads(k, d, hd), _heads(v, d, hd), _heads(iq, d, xd), _dev(ik, d)], dim=0 + ) + ) + + +def _qkv_dense_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0).contiguous() + + +def _qkv_dense_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return _swizzle(torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0)) + + +def _index_qk_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[index-Q | index-K]`` in BF16: this rank's index-Q heads and the replicated index-K, the E4M3 values cast + (the scale tensors, the last two sources, are consumed only).""" + iq, ik = t[0], t[1] + return torch.cat([_heads(iq, d, d.index_dim), _dev(ik, d)], dim=0).to(d.dtype).contiguous() + + +def _gate_up_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return torch.cat([_col(gate, d), _col(up, d)], dim=0).contiguous() + + +def _gate_up_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return _swizzle(torch.cat([_col(gate, d), _col(up, d)], dim=0)) + + +def _row_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _row(t, d).contiguous() + + +def _row_s(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _swizzle(_row(t, d)) + + +def _col_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _col(t, d).contiguous() + + +#: Permute indices of the trtllm-gen expert relayouts, keyed by (shape, family, scale group): one expert's are +#: every expert's, so each is computed once per load. +_PERMUTE_CACHE: dict = {} + + +def _expert_w31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 operand: ``[up ; gate]`` packed-FP4 rows of this rank, then the gated-activation row + reorder and the epilogue-tile shuffle.""" + up, gate = t + w = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + w, _PERMUTE_CACHE, EPILOGUE_TILE_M, is_gated_act_gemm=True + ) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 E4M3 scales: ``[up ; gate]`` rows of this rank, the same row reorder at scale + granularity, then the 128x4 swizzle.""" + up, gate = t + s = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK, is_gated_act_gemm=True + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled.view(torch.uint8).reshape(shape)) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + + +def _expert_w2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 operand: this rank's input columns, then the epilogue-tile row shuffle.""" + w = _row(t, d).contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(w, _PERMUTE_CACHE, EPILOGUE_TILE_M) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 E4M3 scales: this rank's columns, the row shuffle at scale granularity, the swizzle.""" + s = _row(t, d).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w2_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled).view(torch.float8_e4m3fn).reshape(shape) + ) + + +# ------------------------------------------------------------------------------------------------ sources + + +def _one(template: str): + return lambda d, i, t=template: [(_LM + t.format(i=i), None)] + + +def _many(*templates: str): + return lambda d, i, ts=templates: [(tuple(_LM + t.format(i=i) for t in ts), None)] + + +def _per_expert(*suffixes: str): + """One source per routed expert, written at its expert index (all experts are local at tp8).""" + + def build(d, i, sx=suffixes): + out = [] + for e in range(d.num_experts): + keys = tuple(_LM + f"model.layers.{i}.block_sparse_moe.experts.{e}.{s}" for s in sx) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _dense_layers(d): + return range(d.first_sparse) + + +def _sparse_layers(d): + return range(d.first_sparse, d.num_layers) + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense layers, and every layer without the five-way packed + projection (``fuse_qkv_index_projection=False``).""" + return range(d.first_sparse) if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return range(d.first_sparse, d.num_layers) if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else range(d.first_sparse, d.num_layers) + + +_ATTN = "model.layers.{i}.self_attn" +_MOE = "model.layers.{i}.block_sparse_moe" + + +class MiniMaxM3Tp8Weights: + """MiniMax-M3 NVFP4 / sm_107 / tp8: the language model's table and the passes over it. + + ``declare`` allocates every weight it names, ``manifest`` maps checkpoint keys onto them, ``load`` fills + them; ``_targets`` is the one place that decides which parameters exist. + """ + + #: Free the per-expert relayout scratch after each expert stack. + RELEASE_AFTER = "_e_s2" + + WEIGHTS: tuple[W, ...] = ( + # --- the two boundary norms of every layer (Gemma, folded) --- + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + transform=_gemma_fold, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + transform=_gemma_fold, + ), + # --- attention: per-head Q/K norms (Gemma, unfolded), the packed QKV and the output projection --- + W("q_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".q_norm.weight")), + W("k_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".k_norm.weight")), + W( + "qkv_w", + shape=lambda d: (d.qkv_dense_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in "qkv")), + transform=_qkv_dense_w, + layers=_three_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_dense_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many(*(f"{_ATTN}.{p}_proj.weight_scale_inv" for p in "qkv")), + transform=_qkv_dense_s, + layers=_three_way_qkv_layers, + ), + W( + "qkv_w", + shape=lambda d: (d.qkv_sparse_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in ("q", "k", "v", "index_q", "index_k"))), + transform=_qkv_sparse_w, + layers=_five_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_sparse_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + *( + f"{_ATTN}.{p}_proj.weight_scale_inv" + for p in ("q", "k", "v", "index_q", "index_k") + ) + ), + transform=_qkv_sparse_s, + layers=_five_way_qkv_layers, + ), + # The separate BF16 index projection [index-Q | index-K] of the unfused sparse layers. The built-in model + # builds it as an unquantized Linear and fills it through the fused gate/up loader, which casts the + # checkpoint's E4M3 values to BF16 and does not apply their block scales: so does this table (the scale + # tensors are consumed, not applied). + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + src=_many( + _ATTN + ".index_q_proj.weight", + _ATTN + ".index_k_proj.weight", + _ATTN + ".index_q_proj.weight_scale_inv", + _ATTN + ".index_k_proj.weight_scale_inv", + ), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_layers, + ), + W( + "o_w", + shape=lambda d: (d.hidden, d.q_width), + dtype=torch.float8_e4m3fn, + src=_one(_ATTN + ".o_proj.weight"), + transform=_row_w, + ), + W( + "o_s", + shape=lambda d: (d.hidden * (d.q_width // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_ATTN + ".o_proj.weight_scale_inv"), + transform=_row_s, + ), + # --- dense SwiGLU-OAI MLP (layers 0..2) --- + W( + "mlp_gu_w", + shape=lambda d: (2 * d.dense_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight", "model.layers.{i}.mlp.up_proj.weight" + ), + transform=_gate_up_w, + layers=_dense_layers, + ), + W( + "mlp_gu_s", + shape=lambda d: (2 * d.dense_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight_scale_inv", + "model.layers.{i}.mlp.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_dense_layers, + ), + W( + "mlp_dn_w", + shape=lambda d: (d.hidden, d.dense_inter), + dtype=torch.float8_e4m3fn, + src=_one("model.layers.{i}.mlp.down_proj.weight"), + transform=_row_w, + layers=_dense_layers, + ), + W( + "mlp_dn_s", + shape=lambda d: (d.hidden * (d.dense_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one("model.layers.{i}.mlp.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_dense_layers, + ), + # --- MoE (layers 3..59): the fp32 router and its bias, the shared expert --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + dtype=torch.float32, + src=_one(_MOE + ".gate.weight"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "sh_gu_w", + shape=lambda d: (2 * d.shared_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + _MOE + ".shared_experts.gate_proj.weight", _MOE + ".shared_experts.up_proj.weight" + ), + transform=_gate_up_w, + layers=_sparse_layers, + ), + W( + "sh_gu_s", + shape=lambda d: (2 * d.shared_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _MOE + ".shared_experts.gate_proj.weight_scale_inv", + _MOE + ".shared_experts.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_sparse_layers, + ), + W( + "sh_dn_w", + shape=lambda d: (d.hidden, d.shared_inter), + dtype=torch.float8_e4m3fn, + src=_one(_MOE + ".shared_experts.down_proj.weight"), + transform=_row_w, + layers=_sparse_layers, + ), + W( + "sh_dn_s", + shape=lambda d: (d.hidden * (d.shared_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_MOE + ".shared_experts.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_sparse_layers, + ), + # --- the routed experts, NVFP4 (trtllm-gen layouts), all 128 on every rank --- + W( + "e_w31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_per_expert("w3.weight", "w1.weight"), + transform=_expert_w31, + layers=_sparse_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight"), + transform=_expert_w2, + layers=_sparse_layers, + ), + W( + "e_s31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w3.weight_scale", "w1.weight_scale"), + transform=_expert_s31, + layers=_sparse_layers, + ), + W( + "e_s2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w2.weight_scale"), + transform=_expert_s2, + layers=_sparse_layers, + ), + # The per-tensor scalars of every expert, raw; modeling derives the kernels' scalars from them. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + layers=_sparse_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_scale"), + ("isc3", "w3.input_scale"), + ("isc2", "w2.input_scale"), + ("ws2_1", "w1.weight_scale_2"), + ("ws2_3", "w3.weight_scale_2"), + ("ws2_2", "w2.weight_scale_2"), + ) + ), + # --- not per-layer --- + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=_one("model.embed_tokens.weight"), + layers=None, + ), + W( + "final_norm", + shape=lambda d: (d.hidden,), + src=_one("model.norm.weight"), + transform=_gemma_fold, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + sparse = cfg.sparse_attention_config + freq = [int(f) for f in sparse["sparse_attention_freq"]] + tp = mapping.tp_size + head_dim = int(cfg.head_dim) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + index_heads = max(int(sparse["sparse_num_index_heads"]) // tp, 1) + index_dim = int(sparse["sparse_index_dim"]) + return SimpleNamespace( + fused=bool( + getattr( + core.model_config.sparse_attention_config, "fuse_qkv_index_projection", False + ) + ), + num_layers=cfg.num_hidden_layers, + first_sparse=freq.index(1), + hidden=cfg.hidden_size, + head_dim=head_dim, + heads=heads, + kv_heads=kv_heads, + index_heads=index_heads, + index_dim=index_dim, + q_width=heads * head_dim, + qkv_dense_rows=(heads + 2 * kv_heads) * head_dim, + qkv_sparse_rows=(heads + 2 * kv_heads) * head_dim + index_heads * index_dim + index_dim, + dense_inter=int(cfg.dense_intermediate_size) // tp, + shared_inter=int(cfg.shared_intermediate_size) * int(cfg.n_shared_experts) // tp, + moe_inter=int(cfg.intermediate_size) // tp, + num_experts=int(cfg.num_local_experts), + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None)] + return [(f"l{i}_{entry.name}", i) for i in layers] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares (meta under the engine's meta init).""" + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key or key tuple, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, core, weights) -> set: + """The vision tower and its projectors, and the multi-token-prediction head: in the checkpoint, never built + by this target.""" + return {k for k in weights.keys() if k.startswith(_UNBUILT_PREFIXES)} + + def load(self, model, weights) -> None: + """Fill every declared weight and the shell's ``lm_head`` shard; assert coverage both ways.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + assert set(manifest) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest) ^ set(core.w.keys()), + ) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(core, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(core.w[param_key], ckpt, index, transform, weights, d, consumed) + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + # The shell owns the vocab-parallel head: this rank's rows of the full table. + _fill(model.lm_head.weight, _LM + "lm_head.weight", None, _col_w, weights, d, consumed) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + } + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the unbuilt families (vision, MTP head); " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _materialize(entry) -> torch.Tensor: + """``[:]`` on a lazy safetensors slice; a 0-dim tensor (the per-tensor scalars) rejects that index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +def _fill(param, ckpt, index, transform, weights, d, consumed) -> None: + keys = ckpt if isinstance(ckpt, tuple) else (ckpt,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + d.dev = dst.device + # Host tensors: a transform slices this rank's part on the host and moves only that to the device. + src_all = tuple(_materialize(weights[key]) for key in keys) + if transform is not None: + src = transform(src_all if isinstance(ckpt, tuple) else src_all[0], d) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + if src.dim() == 0 and dst.dim() == 0: + pass + elif src.numel() == 1 and dst.numel() == 1: + src = src.reshape(dst.shape) + src = src.to(dst.device, non_blocking=True) + assert tuple(dst.shape) == tuple(src.shape), (ckpt, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + _release(weights, keys) + + +def _release(weights, keys) -> None: + """Drop consumed checkpoint tensors from the engine's dict as the load goes (bounded host memory).""" + if hasattr(weights, "mark_consumed_keys"): + weights.mark_consumed_keys(keys) + else: + for key in keys: + weights.pop(key, None) + + +MODEL_WEIGHTS = MiniMaxM3Tp8Weights() + + +# -------------------------------------------------------------------------------------------- the drafter + + +def _draft_one(template: str): + return lambda d, i, t=template: [(t, None)] + + +def _draft_many(*templates: str): + return lambda d, i, ts=templates: [(tuple(ts), None)] + + +def _draft_fc_norms(d): + return range(d.capture_layers) + + +def _draft_qkv(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0).contiguous() + + +class MiniMaxM3Tp8DraftWeights: + """The one-layer BF16 Eagle3 drafter (``speculative_model``), tp8. + + The fc over the three captured target layers and its per-layer fc_norms, the layer's three norms, its QKV + (input ``[embeds | hidden]``, so twice the hidden width; Q by heads, K/V by KV heads, a KV head on several + ranks when there are fewer than ranks) and o_proj, its gated MLP, the final norm and the full-vocabulary head + (vocab-parallel). The checkpoint's ``embed_tokens`` is the target's table (the drafter reads the target's); it is + a predicted non-load. + """ + + WEIGHTS: tuple[W, ...] = ( + W( + "fc", + shape=lambda d: (d.hidden, d.capture_layers * d.hidden), + src=_draft_one("fc.weight"), + layers=None, + ), + W( + "fc_norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"fc_norm.{i}.weight", None)], + layers=_draft_fc_norms, + ), + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.input_layernorm.weight"), + layers=None, + ), + W( + "hidden_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.hidden_norm.weight"), + layers=None, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.post_attention_layernorm.weight"), + layers=None, + ), + W( + "qkv", + shape=lambda d: (d.qkv_rows, 2 * d.hidden), + src=_draft_many(*(f"layers.0.self_attn.{p}_proj.weight" for p in "qkv")), + transform=_draft_qkv, + layers=None, + ), + W( + "o", + shape=lambda d: (d.hidden, d.q_width), + src=_draft_one("layers.0.self_attn.o_proj.weight"), + transform=_row_w, + layers=None, + ), + W( + "gu", + shape=lambda d: (2 * d.inter, d.hidden), + src=_draft_many("layers.0.mlp.gate_proj.weight", "layers.0.mlp.up_proj.weight"), + transform=_gate_up_w, + layers=None, + ), + W( + "dn", + shape=lambda d: (d.hidden, d.inter), + src=_draft_one("layers.0.mlp.down_proj.weight"), + transform=_row_w, + layers=None, + ), + W("final_norm", shape=lambda d: (d.hidden,), src=_draft_one("norm.weight"), layers=None), + ) + + def dims(self, draft) -> SimpleNamespace: + cfg = draft.draft_config.pretrained_config + mapping = draft.draft_config.mapping + tp = mapping.tp_size + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + return SimpleNamespace( + hidden=cfg.hidden_size, + capture_layers=draft.capture_layers, + head_dim=head_dim, + q_width=heads * head_dim, + qkv_rows=(heads + 2 * kv_heads) * head_dim, + inter=cfg.intermediate_size // tp, + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + if entry.layers is None: + return [(entry.name, None)] + return [(f"{entry.name}{i}", i) for i in entry.layers(d)] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, draft) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + d = self.dims(draft) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, draft, weights) -> set: + """The drafter checkpoint's own embedding: the drafter reads the target's table.""" + return {"embed_tokens.weight"} & set(weights.keys()) + + def load(self, draft, weights) -> None: + d = self.dims(draft) + manifest = self.manifest(draft) + assert set(manifest) == set(draft.w.keys()), set(manifest) ^ set(draft.w.keys()) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(draft, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(draft.w[param_key], ckpt, index, transform, weights, d, consumed) + _fill(draft.lm_head.weight, "lm_head.weight", None, _col_w, weights, d, consumed) + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + f"drafter checkpoint coverage: unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = MiniMaxM3Tp8DraftWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py index 4d5623d7f834..3729e94c9ea9 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 @@ -19,6 +19,7 @@ from __future__ import annotations import json +import os from pathlib import Path from typing import Any, Optional, Tuple @@ -30,10 +31,17 @@ # ``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 -# multicast all-reduce) are built and measured on GB300 only. +# The GPU architectures these targets are written for: the tp4 / tp8 decode +# kernels (the tcgen05 GEMVs, the fused MoE engines, the attention megakernels +# and the multicast all-reduce) are built and measured on GB300 (sm103); the +# tp2 and tp4 Rubin (sm107) targets are ports, switched on one certified path +# at a time. _SM = (10, 3) +_SM_RUBIN = (10, 7) + +# The routing knobs of the former separate Rubin tp2 wide-decode target. Its 9..192-row decode cells are now part of the +# one Rubin tp2 target, so there is nothing left to select: a deployment that still sets either knob is stale. +_REMOVED_TP2_KNOBS = ("TRTLLM_M3_TP2_WIDE", "TRTLLM_M3_TP2_WIDE_MIN_BATCH") # The checkpoint's config shape, read from ``text_config``; see the note in the # gpt_oss routing module for what shape-sniffing does and does not pin. @@ -88,9 +96,9 @@ def _parallel(m) -> Optional[str]: Attention is head-split and the routed experts are tensor-parallel with every expert on every rank: the expert split decides which expert weights each rank loads and into which shapes (per-rank expert width 768 at tp4, - 384 at tp8), so it selects a target rather than a runtime branch. + 384 at tp8, 1536 at tp2), so it selects a target rather than a runtime branch. """ - for tp in (4, 8): + for tp in (2, 4, 8): if ( m.world_size == tp and m.tp_size == tp @@ -271,6 +279,57 @@ def route( return ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + # The same checkpoint on 2 Rubin GPUs (sm107), tensor-parallel attention and + # experts, one-model Eagle3 or no speculative decoding, at any batch size (decode rows 1..192). + if ( + compute_capability == _SM_RUBIN + and shape == _SHAPE + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp2" + and speculation + and drafter is None + and runtime_miss is None + ): + stale = [knob for knob in _REMOVED_TP2_KNOBS if knob in os.environ] + if stale: + raise ValueError( + f"{', '.join(stale)}: removed; the MiniMax-M3 sm107 tp2 target serves every batch size (its " + "wide-decode cells are part of it), so unset these variables" + ) + from .minimax_m3_nvfp4__sm_107__tp2.modeling import ModelingV2MiniMaxM3Nvfp4Sm107Tp2 + + return ModelingV2MiniMaxM3Nvfp4Sm107Tp2 + + # The same checkpoint on 4 Rubin GPUs (sm107), tensor-parallel attention and + # experts, one-model Eagle3 or no speculative decoding. + if ( + compute_capability == _SM_RUBIN + and shape == _SHAPE + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp4" + and speculation + and drafter is None + and runtime_miss is None + ): + from .minimax_m3_nvfp4__sm_107__tp4.modeling import ModelingV2MiniMaxM3Nvfp4Sm107Tp4 + + return ModelingV2MiniMaxM3Nvfp4Sm107Tp4 + + # The same checkpoint on 8 Rubin GPUs over two trays (sm107), tensor-parallel + # attention and experts, one-model Eagle3 or no speculative decoding. + if ( + compute_capability == _SM_RUBIN + and shape == _SHAPE + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp8" + and speculation + and drafter is None + and runtime_miss is None + ): + from .minimax_m3_nvfp4__sm_107__tp8.modeling import ModelingV2MiniMaxM3Nvfp4Sm107Tp8 + + return ModelingV2MiniMaxM3Nvfp4Sm107Tp8 + 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 " @@ -280,8 +339,8 @@ def route( 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 (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" + "minimax_m3 targets are the MiniMax-M3 NVFP4 checkpoint on sm103 (GB300) with 4 or 8 ranks, or on sm107 " + "(Rubin) with 2, 4 or 8 ranks, of tensor-parallel attention and experts, with one-model Eagle3 (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/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 448d64df9925..1f36d0c456ae 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -53,6 +53,18 @@ } ) +# SM107 has no TRTLLM-Gen FMHA, so the thop paged-context kernel serves every +# context phase there, at the 128-token pages MiniMax-M3's KV managers opt in +# (as before those managers listed 128 as a TRTLLM-Gen extra page size); the +# SM107 P128 paged-context oracle covers it. Every other SM keeps the refusal. +_THOP_PAGED_CONTEXT_EXTRA_PAGES_SMS: frozenset = frozenset({107}) + + +def _thop_serves_extra_page_sizes() -> bool: + """Whether the thop paged-context kernel serves a manager's TRTLLM-Gen extra page sizes on this device.""" + return get_sm_version() in _THOP_PAGED_CONTEXT_EXTRA_PAGES_SMS + + # ``thop.attention`` kwargs hard-wired to a literal at the call site (no # rich object owns them). Sync test enforces both the kwarg name and the # literal value. @@ -255,6 +267,7 @@ def _is_supported( and phase != FmhaPhase.GENERATION and metadata.use_paged_context_fmha and metadata.tokens_per_block in extra_page_sizes + and not _thop_serves_extra_page_sizes() ): raise RuntimeError( f"FallbackFmha cannot serve paged prefill with " diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py index e40d8ce0d41c..c49ad2fc908d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py @@ -14,6 +14,8 @@ """Experimental task-scheduled attention entry points.""" +from . import _worktile_scope # noqa: F401 (first: scopes task scheduling's WorkTileInfo rebind) + from .block_sparse import ( BlockSparsePagedTSWrapper, BlockSparseTSWrapper, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py index 1a9d42857f94..d557504e2651 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py @@ -14,6 +14,7 @@ """Compilation and caching for PrimTS block-sparse attention adapters.""" +from .._worktile_scope import with_scalar_work_tile_info from collections.abc import Callable import functools @@ -29,6 +30,7 @@ _COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 3" +@with_scalar_work_tile_info def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: """Compile one contiguous or paged attention adapter for ``key``. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_worktile_scope.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_worktile_scope.py new file mode 100644 index 000000000000..a54e0d6f9893 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_worktile_scope.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Keep CuTe DSL task scheduling's WorkTileInfo rebind inside PrimTS. + +Importing ``cutlass.experimental.task_scheduling`` rebinds six attributes of CUTLASS' +``WorkTileInfo`` (``__init__``, the two MLIR value hooks, ``tile_idx``, ``is_valid_tile`` and +``update_from``) to a scalar three-axis form, for every kernel in the process. Kernels that +subclass ``WorkTileInfo`` with a different tile shape (FA4's four-axis tile) then fail to trace. + +This module is imported first by the package. It snapshots the original attributes, imports +the task scheduling resources (which apply the rebind), snapshots the scalar set, and puts the +originals back. ``scalar_work_tile_info`` installs the scalar set only while a PrimTS kernel +traces and compiles; compiled kernels do not call these Python methods. Nothing about the +generated code changes. Kill switch: TRTLLM_PRIMS_TS_SCOPED_WORKTILE=0 keeps the global rebind. +""" + +import contextlib +import functools +import os +import sys +import threading + +from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo as _WorkTileInfo + +_ATTRS = ( + "__init__", + "__extract_mlir_values__", + "__new_from_mlir_values__", + "tile_idx", + "is_valid_tile", + "update_from", +) +_MISSING = object() +_RESOURCES = "cutlass.experimental.task_scheduling.resources" + + +def _snapshot() -> dict: + return {name: _WorkTileInfo.__dict__.get(name, _MISSING) for name in _ATTRS} + + +def _install(attrs: dict) -> None: + for name, value in attrs.items(): + if value is _MISSING: + if name in _WorkTileInfo.__dict__: + delattr(_WorkTileInfo, name) + else: + setattr(_WorkTileInfo, name, value) + + +# Only scope the rebind when this module saw the originals, i.e. task scheduling was not imported +# before PrimTS. Otherwise keep whatever is installed (the previous global behavior). +_ORIGINAL = None if _RESOURCES in sys.modules else _snapshot() +import cutlass.experimental.task_scheduling.resources # noqa: E402,F401 (applies the scalar rebind) + +_SCALAR = _snapshot() +SCOPED = _ORIGINAL is not None and os.environ.get("TRTLLM_PRIMS_TS_SCOPED_WORKTILE", "1") != "0" +_LOCK = threading.RLock() +_DEPTH = 0 +if SCOPED: + _install(_ORIGINAL) + + +@contextlib.contextmanager +def scalar_work_tile_info(): + """Install the task scheduling WorkTileInfo form while a PrimTS kernel traces (re-entrant).""" + global _DEPTH + if not SCOPED: + yield + return + with _LOCK: + if _DEPTH == 0: + _install(_SCALAR) + _DEPTH += 1 + try: + yield + finally: + _DEPTH -= 1 + if _DEPTH == 0: + _install(_ORIGINAL) + + +def with_scalar_work_tile_info(fn): + """Decorator form of ``scalar_work_tile_info`` for the PrimTS compile functions.""" + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with scalar_work_tile_info(): + return fn(*args, **kwargs) + + return wrapper diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index 1163324ec632..171c02224d12 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -35,6 +35,7 @@ register trace templates. """ +from ._worktile_scope import with_scalar_work_tile_info from dataclasses import dataclass import functools import importlib.metadata @@ -1753,6 +1754,7 @@ def _paged_context_compile_spec( @functools.cache +@with_scalar_work_tile_info def _get_compiled_context( compile_spec: _ContextCompileSpec, ): @@ -1973,6 +1975,7 @@ def fake_compact(dtype, shape, assumed_align): @functools.cache +@with_scalar_work_tile_info def _get_compiled_paged_context( compile_spec: _PagedContextCompileSpec, ): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index 3fcec5efdc2a..a161190007f4 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -15,6 +15,7 @@ """Task-scheduled paged decode with a FlashInfer-style plan/run lifecycle.""" +from ._worktile_scope import with_scalar_work_tile_info from collections.abc import Callable, Sequence from dataclasses import dataclass import functools @@ -1510,6 +1511,7 @@ def _make_decode_compile_spec( @functools.cache +@with_scalar_work_tile_info def _get_compiled_decode( compile_spec: _DecodeCompileSpec, ): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py index 5728a7c6d373..e5f8de86571f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py @@ -26,6 +26,7 @@ summary; no route payload is constructed. """ +from ..._worktile_scope import with_scalar_work_tile_info import functools from collections.abc import Callable @@ -450,6 +451,7 @@ def _fake_compact( @functools.cache +@with_scalar_work_tile_info def compile_block_sparse_inspection( *, device_index: int, @@ -495,6 +497,7 @@ def compile_block_sparse_inspection( @functools.cache +@with_scalar_work_tile_info def compile_paged_block_sparse_metadata_inspection( *, device_index: int, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index 3d5fb159a329..0e82e5330728 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -15,6 +15,7 @@ """Task-scheduled paged MLA decode with a plan/run lifecycle.""" +from ._worktile_scope import with_scalar_work_tile_info from collections.abc import Callable from dataclasses import dataclass, field import functools @@ -1044,6 +1045,7 @@ def _make_mla_decode_compile_spec( @functools.cache +@with_scalar_work_tile_info def _get_compiled_mla_decode( compile_spec: _MLADecodeCompileSpec, ): 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 6182057757a3..ee23ee314816 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -28,6 +28,7 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple import torch @@ -38,6 +39,7 @@ TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor, + get_sm_version, prefer_pinned, ) from tensorrt_llm.bindings import DataType @@ -177,6 +179,11 @@ def set_index_v(self, layer_idx: int, out_cache_loc: torch.Tensor, idx_v: torch. buf.index_copy_(0, out_cache_loc.to(torch.long), idx_v.to(buf.dtype)) +# SM107: price M3's real per-token footprint in the pre-allocation estimate (index-K on every sparse layer, FP8 dense +# and shared Eagle layers) instead of the generic all-NVFP4 K/V sum. Kill switch: False keeps the generic estimate. +_SM107_EXACT_KV_SIZE_ESTIMATE = True + + def shared_draft_layer_count(spec_config, layer_mask) -> int: """How many one-model draft layers the base manager appends to this one. @@ -258,6 +265,11 @@ def extend_attention_op_pools_for_shared_draft_layers( ) +# Zero every KV pool view (NVFP4 + FP8 K/V, NVFP4 block scales, paged index keys) once at allocation, so no reader +# ever sees uninitialized bytes. Kill switch (diagnostics only): TRTLLM_M3_KV_POOL_ZERO_INIT=0. +_KV_POOL_ZERO_INIT = os.environ.get("TRTLLM_M3_KV_POOL_ZERO_INIT", "1") == "1" + + class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): """KVCacheManagerV2 subclass with a V2-managed paged index-K cache per sparse layer. @@ -293,6 +305,7 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): # Decode-step shortcuts a model opts into (RuntimeHooks.msa); see configure_decode_shortcuts. skip_unchanged_block_offsets = False reuse_pool_views = False + block_ids_from_views = False _main_kv_mapper_kind = MapperKind.HND _main_kv_view_layout = "NHD" @@ -436,6 +449,19 @@ def __init__( self._block_scale_buffers[layer_idx, kv_layout] = self.get_block_scale_buffers( layer_idx, kv_layout ) + if _KV_POOL_ZERO_INIT: + # The KV pools are written only at live slots, and readers that consume a whole block or page (the CSR + # sparse prefill's V dequant, the full-context decode on a last page's tail) would turn uninitialized + # bytes into 0 * NaN = NaN. Zero every layer's K/V bytes (NVFP4 sparse, FP8 dense, the drafter's in its + # own manager), the NVFP4 block scales and the paged index keys once, before any request runs: every + # byte a kernel writes is unchanged. + for view in ( + *self._kv_buffers.values(), + *self._block_scale_buffers.values(), + *self._index_k_buffers.values(), + ): + if view is not None: + view.zero_() def _validate_speculative_config(self, spec_config: DecodingBaseConfig | None) -> None: # Dynamic-tree relocation moves packed K/V bytes without their NVFP4 @@ -494,6 +520,72 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): * self.head_dim_per_layer[int(local_layer_idx)] ) + @classmethod + def get_cache_size_per_token(cls, model_config, mapping, num_layers=None, **kwargs): + """Pre-allocation bytes per token, including what the generic estimate leaves out. + + The base classmethod sums K/V for the target layers only, all at the cache dtype. This manager also holds + one index-K head per sparse layer (replicated on every TP rank), keeps the dense target layers and the + appended one-model Eagle layers in FP8 when the cache is NVFP4 (see :meth:`_build_cache_config`), and + allocates those shared draft layers itself. On SM107 the difference is added to the per-token term. + """ + base = super().get_cache_size_per_token(model_config, mapping, num_layers, **kwargs) + if ( + not _SM107_EXACT_KV_SIZE_ESTIMATE + or get_sm_version() != 107 + or kwargs.get("is_draft", False) + ): + return base + extra = cls._extra_bytes_per_token(model_config, mapping, num_layers, kwargs) + if isinstance(base, tuple): + return (base[0] + extra,) + tuple(base[1:]) + return base + extra + + @classmethod + def _extra_bytes_per_token(cls, model_config, mapping, num_layers, kwargs) -> int: + config = model_config.pretrained_config + config = getattr(config, "text_config", None) or config + num_target_layers = int(num_layers if num_layers is not None else config.num_hidden_layers) + local_layers = ( + list(mapping.pp_layers(num_target_layers)) + if mapping.pp_size > 1 + else list(range(num_target_layers)) + ) + sparse_layers = [layer for layer in local_layers if layer >= 3] + dense_layers = [layer for layer in local_layers if layer < 3] + tp = 1 if mapping.enable_attention_dp else mapping.tp_size + kv_heads = max(1, int(config.num_key_value_heads) // tp) + head_dim = int( + getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + ) + + sparse_config = getattr(model_config, "sparse_attention_config", None) + index_dim = int(getattr(sparse_config, "sparse_index_dim", None) or 128) + index_elem = 1 if str(getattr(sparse_config, "indexer_kv_dtype", "bf16")) == "fp8" else 2 + extra = len(sparse_layers) * index_dim * index_elem + + kv_elems = 2 * kv_heads * head_dim # K and V + quant_config = model_config.quant_config + nvfp4 = quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache() + if nvfp4: + # The base prices every layer as NVFP4 data + block scales; dense layers are FP8. + extra += len(dense_layers) * (kv_elems - (kv_elems // 2 + kv_elems // 16)) + draft_layer_bytes = kv_elems + else: + draft_layer_bytes = None + spec_config = kwargs.get("spec_config") + if spec_config is not None and not kwargs.get("use_separate_draft_kv_cache", False): + num_draft = shared_draft_layer_count(spec_config, None) + if draft_layer_bytes is None: + elem = ( + 1 + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache() + else 2 + ) + draft_layer_bytes = kv_elems * elem + extra += num_draft * draft_layer_bytes + return extra + def is_nvfp4_layer(self, layer_idx: int) -> bool: """Whether ``layer_idx`` stores packed NVFP4 data and block scales.""" return self.dtype == DataType.NVFP4 and int(layer_idx) in self.sparse_layer_ids @@ -861,7 +953,10 @@ def get_buffers( ) def configure_decode_shortcuts( - self, skip_unchanged_block_offsets: bool, reuse_pool_views: bool + self, + skip_unchanged_block_offsets: bool, + reuse_pool_views: bool, + block_ids_from_views: bool = False, ) -> None: """Turn on the decode-step shortcuts a model opted into. @@ -869,9 +964,14 @@ def configure_decode_shortcuts( whose inputs equal the previous copy's into the same destination. ``reuse_pool_views``: hand out one pool view per geometry instead of re-wrapping the pool memory on every attention call. + ``block_ids_from_views``: fill :meth:`get_block_ids_per_seq` rows from the + page-index views directly (see there). """ self.skip_unchanged_block_offsets = skip_unchanged_block_offsets self.reuse_pool_views = reuse_pool_views + self.block_ids_from_views = block_ids_from_views + if block_ids_from_views: + logger.info("[M3 host] block-id tables filled from page-index views") def _pool_view( self, @@ -1095,7 +1195,22 @@ def get_block_ids_per_seq(self, request_ids): The rows are written through a numpy view of a single zero-filled, pinned result, so the attention metadata builders ship it to the device in one asynchronous copy. + + With ``block_ids_from_views`` each row is slice-copied from the request's + page-index view (an int32 array) rather than built as a Python list of the + whole padded row first; the table is the same. """ + if self.block_ids_from_views: + views = [self.kv_cache_map[req_id].get_base_page_indices(0) for req_id in request_ids] + max_blocks = max((len(view) for view in views), default=0) + padded_tensor = torch.zeros( + (len(views), max_blocks), dtype=torch.int32, pin_memory=prefer_pinned() + ) + rows = padded_tensor.numpy() + for row, view in zip(rows, views): + row[: len(view)] = view + rows[rows == BAD_PAGE_INDEX] = 0 + return padded_tensor block_ids_per_seq = self.get_batch_cache_indices(request_ids) batch = len(block_ids_per_seq) max_blocks = max((len(block_ids) for block_ids in block_ids_per_seq), default=0) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py index b3cbb2716240..6504f0fbd2d4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py @@ -8,6 +8,39 @@ import torch +# E4M3 casts saturate to +-448, the bits the CUDA FP8 writers produce (__NV_SATFINITE / cvt.rn.satfinite); NaN stays +# NaN. Used by the E4M3 cache writes below and by the targets' transient FP8 query casts. Where torch's own E4M3 cast +# already saturates (c10's float8_e4m3fn conversion, shared by its CPU and CUDA kernels) the cast alone is used: an +# explicit clamp would add a kernel per cast for the same bytes. A torch whose cast does not saturate gets the clamp. +# Kill: False (the plain cast, whatever it does). +FP8_SATURATE = True +_E4M3_MAX = 448.0 + + +def _cast_saturates() -> bool: + """Whether this torch's E4M3 cast maps every out-of-range value, infinities included, to +-448.""" + probe = torch.tensor([449.0, 463.0, 464.0, 1e4, float("inf"), -464.0, -float("inf")]) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0]) + return torch.equal(probe.to(torch.float8_e4m3fn).float(), want) + + +# Probed once at import on the CPU, so no CUDA work runs here (nor later inside a graph capture). +_CAST_SATURATES = _cast_saturates() + + +def fp8_e4m3_saturate(values: torch.Tensor) -> torch.Tensor: + """``values`` as float8_e4m3fn, saturated to +-448 (bitwise the plain cast for every finite in-range value). + Values already in float8_e4m3fn are returned as they are (they are finite by construction or already NaN).""" + if values.dtype == torch.float8_e4m3fn: + return values + if FP8_SATURATE and not _CAST_SATURATES: + values = values.clamp(-_E4M3_MAX, _E4M3_MAX) + return values.to(torch.float8_e4m3fn) + + +def _to_cache_dtype(values: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return fp8_e4m3_saturate(values) if dtype == torch.float8_e4m3fn else values.to(dtype) + def write_kv_slots( cache: torch.Tensor, @@ -60,11 +93,11 @@ def write_kv_slots( if layout == "HND": # Advanced indices on dims 0 and 2 broadcast to [num_tokens] and # move front, giving a [num_tokens, num_heads, channel] target. - cache[page, :, within, :] = values.to(cache.dtype) + cache[page, :, within, :] = _to_cache_dtype(values, cache.dtype) else: - cache[page, within] = values.to(cache.dtype) + cache[page, within] = _to_cache_dtype(values, cache.dtype) else: - cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) + cache.index_copy_(0, out_cache_loc.to(torch.long), _to_cache_dtype(values, cache.dtype)) __all__ = ["write_kv_slots"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py index 071f907067b4..0150347026a1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py @@ -799,8 +799,8 @@ def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> i return 1 << (target.bit_length() - 1) -# NVFP4 instructions are supported on both datacenter Blackwell targets. -_SM100F_CAPABILITIES = ((10, 0), (10, 3)) +# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) # Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. _NVFP4_TUNED_CAPABILITIES = ((10, 3),) @@ -1063,7 +1063,7 @@ def minimax_m3_sparse_attn_decode( capability = torch.cuda.get_device_capability(q.device) if capability not in _SM100F_CAPABILITIES: raise NotImplementedError( - "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103; " + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; " f"got SM{capability[0] * 10 + capability[1]}." ) _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py index 075ebe04da12..ec9fdcd89809 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py @@ -25,6 +25,9 @@ import torch +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers from .msa_utils import check_decode_span_shape @@ -278,7 +281,7 @@ def minimax_m3_trtllm_gen_dense_decode( # q/k/v, so an FP8 pool needs FP8 Q. M3 stores unscaled E4M3, so this is a # plain cast and a no-op when the fused producer already emitted FP8. if kv_pool.dtype == torch.float8_e4m3fn and q.dtype != torch.float8_e4m3fn: - q = q.to(torch.float8_e4m3fn) + q = fp8_e4m3_saturate(q) reserve = torch.cuda.is_current_stream_capturing() workspace_size = _workspace(q.dtype, num_heads, int(q.shape[2]), int(kv_pool.shape[1])) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index dd6d0b5c7436..c1e988803fec 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -39,10 +39,11 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs -from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed, copy_h2d_pool_shared from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy from tensorrt_llm._utils import maybe_pin_memory from tensorrt_llm.bindings import DataType +from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig from .common import ( @@ -111,9 +112,23 @@ def _host_prefix_sum(lens_cpu: torch.Tensor, dtype: torch.dtype) -> torch.Tensor return maybe_pin_memory(torch.cat([sums.new_zeros(1), sums]).to(dtype)) +# Kill switches for A/B only. +# SUBPAGE_IN_GRAPH: the sub-page block table is expanded from msa_block_table on the device inside the forward, once, +# by the first dense layer that reads it (msa_subpage_rows), so the block table can join the step-input pack and the +# step boundary issues no expansion; False expands it in prepare() after an unpacked block-table upload, as before. +SUBPAGE_IN_GRAPH = True +# SKIP_UNCHANGED_BLOCK_TABLE: msa_block_table takes the H2D skip cache. Its only writer is prepare() through that cache +# (the device and every captured graph only read it), so the shadow describes the device rows; False always uploads it. +SKIP_UNCHANGED_BLOCK_TABLE = True +# TAIL_FILL_IN_FORWARD: prepare() records where msa_out_cache_loc's live slots end and the forward invalidates the rest +# (on_update_kv_lens, which opens every forward, so a captured forward fills its own tail at each replay) instead of an +# eager fill between the decode graphs; False fills it in prepare() as before. +TAIL_FILL_IN_FORWARD = True + + def _fill_unused_slots(buf: torch.Tensor, start: int) -> None: - """buf[start:] = -1 after a write of the live slots buf[:start] (see h2d_cache.fill_tail_if_changed).""" - h2d_cache.fill_tail_if_changed(buf, start, -1) + """buf[start:] = -1 after a write of the live slots buf[:start] (see h2d_cache.fill_tail_pool_shared).""" + h2d_cache.fill_tail_pool_shared(buf, start, -1) class MsaDecodeSpan(NamedTuple): @@ -196,6 +211,11 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # factor, or 0 where the pool has no single one; see msa_subpage_rows. msa_subpage_block_table: Optional[torch.Tensor] = None _msa_subpages_per_slot: int = 0 + # SUBPAGE_IN_GRAPH: the rows prepare() staged for the forward to expand (0: none), and whether this forward did. + _msa_subpage_rows: int = 0 + _msa_subpage_expanded: bool = False + # TAIL_FILL_IN_FORWARD: the first msa_out_cache_loc row past this step's live slots (None: prepare() filled it). + _msa_tail_start: Optional[int] = None # Inputs for on_update_kv_lens: each query token's request row and offset # within the request, and the kv_lens prepare() staged (the upper bound the # correction is clamped to). The write slots are re-derived from @@ -370,6 +390,15 @@ def msa_subpage_rows(self, row_first: int, row_last: int) -> Tuple[Optional[torc table = self.msa_subpage_block_table if table is None: return None, 0 + rows = self._msa_subpage_rows + if rows > 0 and not self._msa_subpage_expanded: + # Once per forward (on_update_kv_lens, which opens every forward, resets it), from the block table the + # pack landed, over the rows prepare() staged; a captured forward records it once, so every replay + # expands its own rows. + write_subpage_block_table( + self.msa_block_table[:rows], self._msa_subpages_per_slot, table[:rows] + ) + self._msa_subpage_expanded = True return table[row_first:row_last], self._msa_subpages_per_slot def _msa_main_kv_is_fp8(self) -> bool: @@ -706,9 +735,9 @@ def _stage_msa_short_table(self, n_valid: torch.Tensor, params) -> None: ar = torch.arange(topk, dtype=torch.int32) table = torch.where(ar[None, :] < n_valid.to(torch.int32)[:, None], ar[None, :], -1) table = table.unsqueeze(0).expand(num_kv_heads, total_q, topk).contiguous() - copy_h2d_if_changed( - self.msa_short_table[: table.numel()], maybe_pin_memory(table.flatten()), pack=True - ) + # Uploaded every step, never skipped: the buffer is shared through the CUDA graph buffer pool, so the skip + # cache's host shadow is not proof of its device contents, and a stale table selects too few blocks. + step_inputs.copy_(self.msa_short_table[: table.numel()], maybe_pin_memory(table.flatten())) self._msa_short_ok = True def prepare(self) -> None: @@ -882,6 +911,18 @@ def _msa_kv_lens_may_change(self) -> bool: or (runtime_features is not None and runtime_features.has_speculative_draft_tokens) ) + def _fill_tail_in_forward(self) -> None: + """TAIL_FILL_IN_FORWARD: -1 over msa_out_cache_loc past the live slots prepare() staged (see the prepare() + comment on padded rows), before anything in the forward reads it.""" + start = self._msa_tail_start + buf = self.msa_out_cache_loc + if start is not None and buf is not None and start < buf.shape[0]: + logger.info_once( + "[M3 host] MSA out_cache_loc tail filled inside the forward", + key="msa_tail_fill_in_forward", + ) + buf[start:].fill_(-1) + def on_update_kv_lens(self) -> None: """Re-derive lengths, slots and valid-block counts from the corrected kv_lens_cuda. @@ -901,6 +942,9 @@ def on_update_kv_lens(self) -> None: extend_ctx promotes speculative generation rows into the context prefix. Staged host bounds remain valid upper bounds as lengths shrink. """ + # A new forward: the first dense layer that reads the sub-page table expands it again (msa_subpage_rows). + self._msa_subpage_expanded = False + self._fill_tail_in_forward() super().on_update_kv_lens() if not self._msa_fields_ready or not self._msa_kv_lens_dynamic: return @@ -1117,6 +1161,8 @@ def _build_msa_fields(self) -> None: decoding the inputs for on_update_kv_lens are staged as well. """ self._msa_fields_ready = False + # Set again below once this step's live slots are staged; until then the forward fills nothing. + self._msa_tail_start = None if not self._msa_buffers_ready: return request_ids = self.request_ids @@ -1194,7 +1240,11 @@ def _build_msa_fields(self) -> None: # would make its shape dynamic, and a negative slot is what makes those # rows cache-write no-ops instead. Invalidate the unwritten tail so they # cannot reuse the previous step's live slots, which address real pages. - _fill_unused_slots(self.msa_out_cache_loc, total_new_tokens) + if TAIL_FILL_IN_FORWARD: + self._msa_tail_start = total_new_tokens + else: + self._msa_tail_start = None + _fill_unused_slots(self.msa_out_cache_loc, total_new_tokens) if kv_indices is not None: step_inputs.copy_(self.msa_kv_indices[: int(kv_indices.shape[0])], kv_indices) @@ -1202,12 +1252,22 @@ def _build_msa_fields(self) -> None: # from the same host block ids the flat page table was built from. # Columns past a request's page count are left stale rather than # cleared: every consumer bounds its walk by seq_lens. - block_table_changed = copy_h2d_if_changed( - self.msa_block_table[:batch_size, :block_table_cols], - maybe_pin_memory(block_ids_cpu.to(torch.int32)), - # write_subpage_block_table below reads it on the device before the replay. - pack=self.msa_subpage_block_table is None, - ) + # With SUBPAGE_IN_GRAPH nothing reads the table on the device before the replay, so it can join the pack; + # the prepare-time expansion below reads it first. + block_table_host = maybe_pin_memory(block_ids_cpu.to(torch.int32)) + pack_block_table = self.msa_subpage_block_table is None or SUBPAGE_IN_GRAPH + if SKIP_UNCHANGED_BLOCK_TABLE: + block_table_changed = copy_h2d_if_changed( + self.msa_block_table[:batch_size, :block_table_cols], + block_table_host, + pack=pack_block_table, + ) + else: + block_table_changed = copy_h2d_pool_shared( + self.msa_block_table[:batch_size, :block_table_cols], + block_table_host, + pack=pack_block_table, + ) # Packed only when the device scan below (h2d_cache off) does not read it before the replay. if h2d_cache.enabled(): step_inputs.copy_(self.msa_seq_lens_cuda[:batch_size], kv_lens_cpu) @@ -1219,8 +1279,8 @@ def _build_msa_fields(self) -> None: if h2d_cache.enabled(): # Both prefix sums from the host lengths (the values the device scans would give), staged with # their leading zero: copies instead of zero / scan kernels between the decode graphs. - copy_h2d_if_changed(qo_lens_cuda, qo_lens_cpu, pack=True) - copy_h2d_if_changed( + copy_h2d_pool_shared(qo_lens_cuda, qo_lens_cpu, pack=True) + copy_h2d_pool_shared( self.msa_cu_q_lens[: batch_size + 1], _host_prefix_sum(qo_lens_cpu, self.msa_cu_q_lens.dtype), pack=True, @@ -1278,7 +1338,14 @@ def _build_msa_fields(self) -> None: # Sub-page expansion for the trtllm-gen dense layers, staged once here # instead of once per layer, outside capture into a graph-stable # buffer as with the slot table above. - if self.msa_subpage_block_table is not None and block_table_changed: + # With SUBPAGE_IN_GRAPH the forward expands this step's rows instead (msa_subpage_rows). + self._msa_subpage_expanded = False + self._msa_subpage_rows = batch_size if SUBPAGE_IN_GRAPH else 0 + if ( + not SUBPAGE_IN_GRAPH + and self.msa_subpage_block_table is not None + and block_table_changed + ): write_subpage_block_table( self.msa_block_table[:batch_size], self._msa_subpages_per_slot, @@ -1302,10 +1369,10 @@ def _build_msa_fields(self) -> None: torch.arange(total_new_tokens, dtype=torch.int64) - torch.repeat_interleave(starts, qo_long) ).to(torch.int32) - copy_h2d_if_changed( + copy_h2d_pool_shared( self.msa_q_batch_row[:total_new_tokens], maybe_pin_memory(batch_row_cpu), pack=True ) - copy_h2d_if_changed( + copy_h2d_pool_shared( self.msa_q_intra[:total_new_tokens], maybe_pin_memory(intra_cpu), pack=True ) # The staged lens are the upper bound on_update_kv_lens clamps to: the diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 7a88db7bddd5..abe67148247e 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -41,7 +41,7 @@ from ...pyexecutor import step_inputs from ...pyexecutor.config_utils import is_mla -from ...pyexecutor.h2d_cache import copy_h2d_if_changed +from ...pyexecutor.h2d_cache import copy_h2d_pool_shared from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) from .fmha.manager import FmhaManager @@ -57,6 +57,11 @@ from .sparse.skip_softmax import SkipSoftmaxParams from .utils import log_attention_failure_context +# With an H2D skip cache active (RuntimeHooks.skip_unchanged_h2d), the prompt-length upload of prepare() joins the +# open step-input pack instead of going out as its own copy on the main stream between the step graphs; False +# issues it directly. +PACK_PROMPT_LENS = True + _SKIP_CORRECTION_SUPPORTED_SMS = frozenset((100, 103)) @@ -1008,8 +1013,9 @@ def prepare(self) -> None: device='cpu', ) self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens) - copy_h2d_if_changed(self.prompt_lens_cuda[:self.num_seqs], - self.prompt_lens_cpu[:self.num_seqs]) + copy_h2d_pool_shared(self.prompt_lens_cuda[:self.num_seqs], + self.prompt_lens_cpu[:self.num_seqs], + pack=PACK_PROMPT_LENS) # number of tokens in the kv cache for each sequence in the batch cached_token_lens = torch.tensor( @@ -1170,7 +1176,7 @@ def prepare_encoder_decoder_from_precomputed_lengths( f"({self.kv_cache_manager.max_seq_len}).") num_seqs = self.num_seqs - copy_h2d_if_changed(self.prompt_lens_cuda[:num_seqs], prompt_lens) + copy_h2d_pool_shared(self.prompt_lens_cuda[:num_seqs], prompt_lens) self.kv_lens_cuda[:num_seqs].copy_(kv_lens, non_blocking=True) self.host_total_kv_lens[0] = context_kv_tokens self.host_total_kv_lens[1] = generation_kv_tokens diff --git a/tensorrt_llm/_torch/custom_ops/mxfp8_cute_pinned_tactics.py b/tensorrt_llm/_torch/custom_ops/mxfp8_cute_pinned_tactics.py new file mode 100644 index 000000000000..1511adaed71a --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/mxfp8_cute_pinned_tactics.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pinned FlashInfer CuTe DSL MXFP8 GEMM tactics for SM107 (MiniMax-M3 TP2 projections), above 32 rows. + +TRT-LLM runs FlashInfer's CuTe DSL MXFP8 GEMM with FlashInfer's config sweep skipped, so above 32 rows (the +split-K low-M kernel's limit) every call takes the default tile (128, 128), cluster (1, 1): ceil(N / 128) CTAs, a +quarter of the SMs streaming the weight. This table holds FlashInfer's own tuned tactic per (N, K, M bucket) on +SM107. Every entry is non-split-K (split_k 1, each output keeps its full-K accumulation order), so its output equals +the default tactic's bit for bit (measured on all entries), at 18..40% less time from 33 to 16384 rows. A fixed +table rather than live tuning keeps the choice identical across runs and nodes. Used where TRT-LLM calls the +FlashInfer runner (decode-graph capture, autotune); eager prefill takes other GEMMs. Off unless armed: the executor +calls arm() for a model whose RuntimeHooks set mxfp8_pinned_tactics. Kill: TRTLLM_MXFP8_CUTE_PIN=0. +""" + +import os +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger + +ENABLED = os.environ.get("TRTLLM_MXFP8_CUTE_PIN", "1") == "1" +MIN_ROWS = 33 # 1..32 rows keep FlashInfer's split-K low-M path (its tuned tactics there are split-K: not bitwise) +MAX_BUCKET = 16384 + +Tactic = Tuple[Tuple[int, int], Tuple[int, int], bool, bool, int] +# (N, K, M bucket = next power of two >= M, 64..16384) +# -> (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch, split_k) +TACTICS: Dict[Tuple[int, int, int], Tactic] = { + (3072, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (3072, 6144, 128): ((256, 64), (2, 1), True, False, 1), + (3072, 6144, 256): ((256, 64), (2, 1), True, False, 1), + (3072, 6144, 512): ((256, 64), (2, 1), False, False, 1), + (3072, 6144, 1024): ((256, 128), (2, 1), False, False, 1), + (3072, 6144, 2048): ((256, 256), (2, 1), False, False, 1), + (3072, 6144, 4096): ((256, 192), (4, 2), False, False, 1), + (3072, 6144, 8192): ((256, 192), (2, 1), False, False, 1), + (3072, 6144, 16384): ((256, 192), (2, 1), True, False, 1), + (4608, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (4608, 6144, 128): ((256, 64), (2, 1), True, False, 1), + (4608, 6144, 256): ((256, 64), (2, 4), False, True, 1), + (4608, 6144, 512): ((256, 128), (2, 2), False, False, 1), + (4608, 6144, 1024): ((256, 192), (2, 1), False, False, 1), + (4608, 6144, 2048): ((256, 192), (2, 1), False, False, 1), + (4608, 6144, 4096): ((256, 256), (2, 1), False, False, 1), + (4608, 6144, 8192): ((256, 192), (2, 4), True, False, 1), + (4608, 6144, 16384): ((256, 256), (2, 1), True, False, 1), + (4992, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (4992, 6144, 128): ((256, 64), (2, 1), True, False, 1), + (4992, 6144, 256): ((256, 64), (2, 4), False, True, 1), + (4992, 6144, 512): ((256, 128), (4, 2), False, True, 1), + (4992, 6144, 1024): ((256, 192), (2, 1), False, False, 1), + (4992, 6144, 2048): ((256, 192), (2, 1), False, False, 1), + (4992, 6144, 4096): ((256, 192), (2, 1), False, False, 1), + (4992, 6144, 8192): ((256, 192), (2, 1), False, False, 1), + (4992, 6144, 16384): ((256, 192), (2, 1), True, False, 1), + (6144, 1536, 64): ((256, 64), (2, 1), True, False, 1), + (6144, 1536, 128): ((256, 64), (2, 1), True, False, 1), + (6144, 1536, 256): ((256, 64), (2, 1), True, False, 1), + (6144, 1536, 512): ((256, 128), (2, 1), False, False, 1), + (6144, 1536, 1024): ((256, 256), (2, 1), False, False, 1), + (6144, 1536, 2048): ((256, 192), (2, 4), False, False, 1), + (6144, 1536, 4096): ((256, 192), (2, 1), False, False, 1), + (6144, 1536, 8192): ((256, 192), (2, 1), False, False, 1), + (6144, 1536, 16384): ((256, 192), (2, 1), True, False, 1), + (6144, 4096, 64): ((256, 64), (2, 1), True, False, 1), + (6144, 4096, 128): ((256, 64), (4, 1), True, False, 1), + (6144, 4096, 256): ((256, 64), (2, 1), True, False, 1), + (6144, 4096, 512): ((256, 128), (2, 1), False, False, 1), + (6144, 4096, 1024): ((256, 256), (2, 1), False, False, 1), + (6144, 4096, 2048): ((256, 192), (2, 4), False, False, 1), + (6144, 4096, 4096): ((256, 192), (2, 1), False, False, 1), + (6144, 4096, 8192): ((256, 192), (2, 1), True, False, 1), + (6144, 4096, 16384): ((256, 192), (2, 1), True, False, 1), + (6144, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (6144, 6144, 128): ((256, 64), (2, 1), True, False, 1), + (6144, 6144, 256): ((256, 64), (2, 1), True, False, 1), + (6144, 6144, 512): ((256, 128), (2, 1), True, False, 1), + (6144, 6144, 1024): ((256, 256), (2, 1), False, False, 1), + (6144, 6144, 2048): ((256, 192), (2, 4), False, False, 1), + (6144, 6144, 4096): ((256, 192), (2, 1), False, False, 1), + (6144, 6144, 8192): ((256, 192), (2, 1), False, False, 1), + (6144, 6144, 16384): ((256, 256), (2, 1), True, False, 1), + (12288, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (12288, 6144, 128): ((128, 64), (1, 2), True, False, 1), + (12288, 6144, 256): ((256, 128), (2, 1), False, False, 1), + (12288, 6144, 512): ((256, 256), (2, 1), False, False, 1), + (12288, 6144, 1024): ((256, 192), (2, 2), False, False, 1), + (12288, 6144, 2048): ((256, 192), (2, 1), False, False, 1), + (12288, 6144, 4096): ((256, 192), (2, 1), False, False, 1), + (12288, 6144, 8192): ((256, 256), (2, 1), False, False, 1), + (12288, 6144, 16384): ((256, 256), (2, 1), True, False, 1), +} + +_STATE: dict = {} +# Process-wide once armed: another model later or concurrently in the same process, at a table shape on SM107, +# also takes the pinned tactic. Its outputs are unchanged (each entry is bitwise the default tactic), only faster. +# The LLM API runs one executor per worker process. +_ARMED = [False] + + +def arm() -> None: + """Engage the table for this process (a model opted in through RuntimeHooks.mxfp8_pinned_tactics).""" + _ARMED[0] = True + + +def _bucket(m: int) -> int: + return max(64, 1 << (m - 1).bit_length()) + + +def tactic_for(m: int, n: int, k: int) -> Optional[Tactic]: + """The pinned tactic of an [m, k] x [n, k]^T call, or None (FlashInfer's own path).""" + if not ENABLED or not _ARMED[0] or m < MIN_ROWS or m > MAX_BUCKET or get_sm_version() != 107: + return None + return TACTICS.get((n, k, _bucket(m))) + + +def run( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> Optional[torch.Tensor]: + """The MXFP8 GEMM act @ weight^T on FlashInfer's CuTe DSL runner at the pinned tactic, or None when no entry covers + the call (the caller runs FlashInfer's default path). Bitwise equal to that default path where an entry exists.""" + m, k = act.shape + n = weight.shape[0] + tactic = tactic_for(m, n, k) + if tactic is None: + return None + try: + from flashinfer.gemm import gemm_base as _gb + + key = (out_dtype, act.device) + st = _STATE.get(key) + if st is None: + major, minor = _gb.get_compute_capability(act.device) + st = _STATE[key] = ( + _gb._cute_dsl_gemm_mxfp8_runner(major, minor, True, out_dtype), + _gb._get_cache_buf("mm_mxfp8_workspace", _gb.DEFAULT_WORKSPACE_SIZE, act.device), + ) + logger.info( + "[MXFP8 GEMM] pinned CuTe DSL tactics engaged above 32 rows (TRTLLM_MXFP8_CUTE_PIN)" + ) + runner, workspace = st + out = torch.empty((m, n), dtype=out_dtype, device=act.device) + runner( + inputs=[act, weight.t(), act_scale, weight_scale, out_dtype, out, workspace], + tactic=tactic, + ) + return out + except Exception as e: # noqa: BLE001 - a FlashInfer API change falls back to the default path, once loudly + if "failed" not in _STATE: + _STATE["failed"] = True + logger.warning( + f"[MXFP8 GEMM] pinned CuTe DSL tactic unavailable ({type(e).__name__}: {e}); default path" + ) + return None diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index e7724b6c601f..014085c635a2 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -40,6 +40,7 @@ from ..cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl +from . import mxfp8_cute_pinned_tactics as _mxfp8_cute_pinned from .fast_custom_op import fast_custom_op if IS_FLASHINFER_AVAILABLE: @@ -807,6 +808,13 @@ def forward( backend = "cutlass" context = nullcontext() else: + # Above 32 rows on SM107, for a model that opts in: the pinned tuned tactic (bitwise the + # heuristic's output, faster); see mxfp8_cute_pinned_tactics. Otherwise, or with no entry: + # FlashInfer's heuristic config. + pinned = _mxfp8_cute_pinned.run(act, act_scale, weight, + weight_scale, self.output_dtype) + if pinned is not None: + return pinned # Skip FlashInfer's CuTeDSL config sweep; use its heuristic config. backend = "cute-dsl" context = _flashinfer_autotune(tune_mode=False, @@ -2691,6 +2699,16 @@ def _( return gate.new_empty(gate.shape, dtype=o_dtype) +# The AllReduce tuner profiles NCCL_SYMMETRIC by pre-allocating two NCCL window buffers of [2 x max_num_tokens, hidden] +# (plus registration) inside the memory-profiling warm-up, so they shrink the KV pool even where NCCL_SYMMETRIC never +# wins (agg TP2: only NCCL / ONESHOT are picked; ~1.24 GiB per rank of KV, GEMM-AR 803388). +# TRTLLM_AR_TUNE_NCCL_SYMMETRIC=0 leaves it out of the candidates (no windows) and makes plain NCCL the cache-miss +# fallback; the default keeps stock behavior. +_AR_TUNE_NCCL_SYMMETRIC = os.environ.get("TRTLLM_AR_TUNE_NCCL_SYMMETRIC", + "1") == "1" +_AR_TUNE_NCCL_SYMMETRIC_LOGGED = False + + class AllReduceRunner(TunableRunner): _prealloc_lock: ClassVar[threading.Lock] = threading.Lock() _prealloc_done: ClassVar[set] = set() @@ -2796,10 +2814,16 @@ def get_valid_tactics( profile: OptimizationProfile, **kwargs, ) -> List[int]: - valid_strategies = [ - AllReduceStrategy.NCCL_SYMMETRIC.value, - AllReduceStrategy.NCCL.value, - ] + global _AR_TUNE_NCCL_SYMMETRIC_LOGGED + valid_strategies = [AllReduceStrategy.NCCL.value] + if _AR_TUNE_NCCL_SYMMETRIC: + valid_strategies.insert(0, AllReduceStrategy.NCCL_SYMMETRIC.value) + elif not _AR_TUNE_NCCL_SYMMETRIC_LOGGED: + _AR_TUNE_NCCL_SYMMETRIC_LOGGED = True + logger.info( + "[tunable_allreduce] AR tuner: NCCL_SYMMETRIC left out (no NCCL window buffers " + "pre-allocated); cache-miss fallback is NCCL (TRTLLM_AR_TUNE_NCCL_SYMMETRIC=0)" + ) # Fallback in allreduceOp is set to NCCL_SYMMETRIC as default # So we need to check if the workspace size is too large to avoid hanging. workspace_size = inputs[0].numel() * inputs[0].element_size() @@ -2890,7 +2914,8 @@ def forward( # Fall back to NCCL_SYMMETRIC; asymmetric ncclMemAlloc failures are handled # by a cross-rank barrier in NCCLWindowAllocator which falls back # to plain NCCL. - tactic = AllReduceStrategy.NCCL_SYMMETRIC.value + tactic = (AllReduceStrategy.NCCL_SYMMETRIC.value if + _AR_TUNE_NCCL_SYMMETRIC else AllReduceStrategy.NCCL.value) return torch.ops.trtllm.allreduce( input, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py new file mode 100644 index 000000000000..d6a67225eb01 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuTe DSL decode kernels of the MiniMax-M3 NVFP4 / sm_103 / tp4 modeling_v2 target.""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py new file mode 100644 index 000000000000..4c6fc92cd49c --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 decode embedding gather + first-layer RMSNorm as one CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py new file mode 100644 index 000000000000..acf27a95316b --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. + +128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original +K3 wrapper this target does not promise FlashInfer's bit-identical reduction. +""" + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +NORM_THREADS = 128 +NORM_VEC = 8 +NORM_WARPS = 4 + + +@cute.kernel +def m3_embed_norm_kernel( + ids: cute.Tensor, # int32 or int64 [N] + table: cute.Tensor, # bf16 [V, H] + weight: cute.Tensor, # bf16 [H] + raw: cute.Tensor, # bf16 [N, H], out: the rows + out: cute.Tensor, # bf16 [N, H], out: the normed rows + vocab: cutlass.Int64, + eps: cutlass.Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + hidden: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 + ) + + tok = cutlass.Int64(ids[bidx]) + valid = (tok >= cutlass.Int64(0)) & (tok < vocab) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + + gX = cute.local_tile(table, tiler_mn, (row, 0)) + gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) + gY = cute.local_tile(out, tiler_mn, (bidx, 0)) + w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 + ) + thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + tXgR = thr_copy_O.partition_D(gR) + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. + tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) + if valid: + cute.copy(copy_atom_load, tXgX, tXrX) + cute.copy(copy_atom_load, tWgW, tWrW) + cute.copy(copy_atom_store, tXrX, tXgR) + + # flashinfer's RMSNormKernel arithmetic, in its order (see the header). + x = tXrX.load().to(cutlass.Float32) + x_sq = x * x + sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) + for i in cutlass.range_constexpr(5): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + if lane == 0: + reduction_buffer[0, warp] = sum_sq + cute.arch.barrier() + total = cutlass.Float32(0.0) + if lane < NORM_WARPS: + total = reduction_buffer[0, lane] + for i in cutlass.range_constexpr(5): + total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) + mean_sq = total / cutlass.Float32(hidden) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + w = tXrW.load().to(cutlass.Float32) + y = x * rstd * (w + cutlass.Float32(0.0)) + tXrO.store(y.to(out.element_type)) + cute.copy(copy_atom_store, tXrO, tXgO) + + +@cute.jit +def m3_embed_norm( + ids: cute.Tensor, + table: cute.Tensor, + weight: cute.Tensor, + raw: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int64, + eps: cutlass.Float32, + n_tokens: cutlass.Constexpr[int], + hidden: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. + blocks = hidden // (NORM_VEC * NORM_THREADS) + tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), + stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip + m3_embed_norm_kernel( + ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden + ).launch( + grid=[n_tokens, 1, 1], + block=[NORM_THREADS, 1, 1], + smem=NORM_WARPS * 4, + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py new file mode 100644 index 000000000000..0acfdecb196e --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host side of the MiniMax-M3 decode embedding + first-layer RMSNorm (m3_embedding_firstnorm_kernel.py): H 6144, +1..8 rows, into caller-owned buffers. + +Compiled on the first call per (rows, id dtype, table shape, PDL), which must happen outside CUDA-graph capture +(:class:`DecodeEmbedding` compiles at construction). +""" + +import threading +from typing import Dict, Tuple + +import cuda.bindings.driver as cuda_driver +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit + +HIDDEN = 6144 +MAX_ROWS = 8 + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def arg(tensor): + return from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _check(ids, table, weight, raw, out) -> None: + rows = ids.shape[0] if ids.dim() == 1 else 0 + if not ( + 1 <= rows <= MAX_ROWS + and ids.dtype in (torch.int32, torch.int64) + and ids.is_contiguous() + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN,) + and all( + t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() + for t in (raw, out) + ) + ): + raise ValueError( + "m3_embed_norm: expects ids int32/int64 [1..8], a contiguous bf16 table [V, 6144], bf16 weight [6144] " + "and contiguous bf16 raw / out [rows, 6144]" + ) + + +def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) + fn = _compiled[key] = cute.compile( + _m3_embed_norm_jit, + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + _check(ids, table, weight, raw, out) + fn = _compile(ids, table, weight, eps, raw, out, use_pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return out, raw + + +class DecodeEmbedding: + """Caller-owned buffers of one (rows, id dtype) decode embedding: ``raw`` (the gathered rows, the first residual) + and ``output`` (their RMSNorm), over ``table`` and the first layer's input ``norm``. Built (and compiled) before + CUDA-graph capture.""" + + def __init__(self, table, norm, rows, id_dtype, use_pdl=True): + assert not torch.cuda.is_current_stream_capturing() + assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN + assert ( + norm.weight.dtype == torch.bfloat16 + and norm.weight.shape == (HIDDEN,) + and not norm.use_gemma + ) + assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.norm, self.rows, self.id_dtype, self.use_pdl = ( + table, + norm, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(norm.variance_epsilon) + self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _compile(ids, table, norm.weight, self.eps, self.raw, self.output, use_pdl) + + def __call__(self, ids): + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return m3_embed_norm( + ids, self.table, self.norm.weight, self.eps, self.raw, self.output, self.use_pdl + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py new file mode 100644 index 000000000000..0aa75c0bf82c --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 MoE layers for low-batch decode as CuTe DSL kernels: ``m3_moe`` for 1..8 tokens (op.py) and ``m3_moe16`` +for 9..16 tokens (op16.py).""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py new file mode 100644 index 000000000000..e317e5e3087a --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py @@ -0,0 +1,3871 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 routed + shared experts for 9..16 decode tokens on one TP rank: ``m3_moe16`` (v2). + +Same math as the 1..8-token engine (m3_moe_kernel.py), operation for operation for the routed part (bitwise equal to +TRTLLM-Gen's fused MoE on the fp32 router-GEMV logits); a different schedule for 9..16 rows: + +- Front: the hidden states land in smem by bulk copies in K chunks (the idle ring); all 16 warps run the router + GEMV (m3_router_gemv's per-thread order) and the NVFP4 quantization of B1 from smem; the logits cross CTAs as + tagged words; top-k one warp per token. +- Ring: FC1 / FC2 weight stages by TMA; their scale words go by 4-byte cp.async straight into the instruction atom + layout (SFA_CP; no raw scale box, no splicer warp), completing on the stage's full barrier. +- FC1 -> FC2 hand-off (B2IMG): the FC1 epilogue writes each unit's FP4 values and scale byte into a per-group global + image laid out as the FC2 B tile + its scale atoms; one counter per group (release add per finished unit, every + counter gains exactly 192 per call); warp 5 loads a completed image with two bulk copies. +- Shared expert on NSH more CTAs (MXFP8 tcgen05), as the 1..8 engine. +Configuration via M3_CONFIG (m_max, num_ctas, stages, nb2_cap, lag, shared, sfa_cp, b2buf, ...). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +H = 6144 +I_TP = 768 +TWO_I = 2 * I_TP +E_LOCAL = 128 +TOP_K = 4 +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 12) +assert 9 <= M_MAX <= 16, M_MAX +NCTA = _cfg("num_ctas", 128) +FUSED_ROUTER = True +NROUTE = E_LOCAL +NGEMV = (NROUTE + NCTA - 1) // NCTA +assert NGEMV == 1, "m3_moe16: one router expert per CTA (num_ctas >= 128)" +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +LAG = _cfg("lag", 3) +G_CAP = TOP_K * M_MAX +NG_CAP = M_MAX +THREADS = 512 +NWARPS = THREADS // 32 +GEMV_THREADS = 256 +SHARED = bool(_cfg("shared", 0)) +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +SFA_CP = bool( + _cfg("sfa_cp", 1) +) # A scales by 4-byte cp.async into the atom layout (else raw box + splicer warp) +B2IMG = True +PERM = bool(_cfg("perm", 1)) +PERM_MIN_NG = max( + 3, _cfg("perm_min_ng", 5) +) # (measured: helps at random routing G >= 23, hurts at spec G 13-14) +# a partial last slot group becomes group 1 (the plan must match: op16) +DBG_DELAY = bool( + _cfg("dbg_delay", 0) +) # (stress builds) pseudo-random producer / epilogue / release stalls +DBG_NOSTORE = bool(_cfg("dbg_nostore", 0)) # (timing experiment only: wrong results) +EPI_FENCE = bool( + _cfg("epi_fence", 0) +) # (rel_warp 0) writer lanes' fence.acq_rel.gpu before the release add +EPI_PFENCE = bool( + _cfg("epi_pfence", 1) +) # producer-side fence.proxy.async.global before the release add +REL_WARP = bool( + _cfg("rel_warp", 1) +) # warp 8 does the per-group release adds (off the epilogue's path) +NSH = 24 if SHARED else 0 +SH_I = 768 +SH_KB = H // 32 +SH_DKB = SH_I // 32 +SH_SPLIT = 2 +SH_GT = SH_KB // 4 // SH_SPLIT +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * (SH_I // 128) +SHX_BAR = 6 +SH_PRETILE = False +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = _cfg("n1", 48 if _cfg("m_max", 12) <= 12 else 64) # B rows QS q + t +N1_LD = 1 << (N1 - 1).bit_length() # tcgen05.ld .num (a power of two) covering N1 columns +QS = N1 // 4 +assert QS >= M_MAX +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 +K1_TILES = H // 256 +S1 = K1_TILES // C1 +NUM_BYTES_A = MMA_M * KT_BYTES +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 +KT2 = I_TP // 256 +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = _cfg("n2", (4 * M_MAX + 15) // 16 * 16) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +IMG_B2 = KT2 * NUM_BYTES_B2 # one group's B2 tile image (global and smem) +IMG_SF = KT2 * NUM_KB * 512 # its scale atoms +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) # every group counter gains 192 per call +B1IMG_OFF = None # set below (after N1): the activation image after the NG_CAP group images +TK_WARPS = [0, 1, 2, 3, 6, 8, 9, 10, 11, 12, 13, 14, 5, 7, 15, 4][ + : _cfg("m_max", 12) +] # top-k warp of token t (4 last: early0) +B1_WARP = 15 if _cfg("m_max", 12) <= 14 else -1 # (-1: warp 0 after the routing) +NB2_CAP = _cfg("nb2_cap", 2) +B1IMG_OFF = M_MAX * IMG_BYTES +B1IMG_BYTES = (S1 * N1 * 128 + M_MAX * (H // 16) + 15) // 16 * 16 +IMG_TOTAL = B1IMG_OFF + B1IMG_BYTES +UNITS_ROW = TWO_I // R1 +PARTIAL_EMPTY = -1 +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +TOPK_BAR = 4 +ALL_BAR = 5 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 16) +LG_REPL = _cfg("lg_repl", 64) +LG_STRIDE = M_MAX * E_LOCAL + 32 +W_WARP = 4 +X_WARP = 5 # B2 image loader; slots + plan +S_WARP = 6 # (SFA_CP = 0) scale splicer +MMA_WARP = 7 +EARLY0 = bool(_cfg("early0", 1)) +EVICT_FIRST = 0x12F0000000000000 +EVICT_LAST = 0x14F0000000000000 +_LOG2E = 1.4426950408889634 +NB2BUF = _cfg("b2buf", 2) +SPB = 1 +# x staging: K chunks of XCK elements, NXB smem buffers (rotating if fewer than the chunks) +XCK = _cfg("xck", 2048) +NXC = H // XCK +XCB = M_MAX * XCK * 2 # bytes of one chunk (all tokens) +STAGE_A = NUM_BYTES_A +STAGE_SF = NUM_BYTES_SFA if SFA_CP else NUM_BYTES_SFA_RAW + NUM_BYTES_SFA +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + NB2BUF * IMG_B2 + + IMG_SF + + NB2_CAP * M_MAX * TOP_K * R2 * 2 + + SPB * M_MAX * C1 * R1 * 4 + + (G_CAP + 1) * (NG_CAP + 1) * 4 + + 6 * 1024 +) +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // (STAGE_A + STAGE_SF) +ACC1_COL = 0 +ACC2_COL = 2 * N1 +ACC2_STRIDE = 1 << (N2 - 1).bit_length() +SFB1_COL = ACC2_COL + 2 * ACC2_STRIDE +SFB1_WSTEP = 1 +SFB2_COL = ( + SFB1_COL + NUM_SFB1_ATOMS * 4 +) # one TMEM slot (loaded after the previous group's FC2 MMAs) +SFA1_COL = SFB2_COL + KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# smem: [ring A | ring SF | B2 bufs | B2 SF staging | B1]; x chunks alias [0, _OFF_B1) during the front +_OFF_SFA = STAGES * NUM_BYTES_A +_OFF_B2 = (_OFF_SFA + STAGES * STAGE_SF + 1023) // 1024 * 1024 +_OFF_SFB2 = _OFF_B2 + NB2BUF * IMG_B2 +_OFF_B1 = (_OFF_SFB2 + IMG_SF + 1023) // 1024 * 1024 +BIG = _OFF_B1 + S1 * NUM_BYTES_B1 +NXB = min(NXC, _OFF_B1 // XCB) +assert NXB >= 2, (NXB, _OFF_B1, XCB) +# shared CTA layout in BIG: one pass over all M tokens; B rows SHN q + t (gate_up, 2 K slices), t (down) +SHN = 16 +SH_NPASS = 1 +SH_HINT = EVICT_FIRST +SH_PF_POST = bool( + _cfg("sh_pf_post", 1) +) # L2 prefetch of the shared weights after the grid dependency (else before) +_SB_A = 0 +_SH_FIX = ( + SH_GT * 2 * SHN * 128 + + 6 * SHN * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + 64 * SHN * 4 + + 2 * 32 * SHN * 4 + + 2048 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * 2 * SHN * 128 +_SB_XC = _SB_BH + 6 * SHN * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + 64 * SHN * 4 +_SB_END = _SB_GU + 2 * 32 * SHN * 4 +SH_ACC_GU = 0 +SH_ACC_DN = 2 * SHN +SH_SFA = SH_ACC_DN + 2 * SHN +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +SH_TMEM = 256 +if SHARED: + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) +TD_STRIDE = 64 +PFD = _cfg("pfd", 0) # L2 prefetch distance in ring stages (0: off) +PF_WARP = 12 +TICKET = bool(_cfg("ticket", 0)) # CTA roles by entry order +# priority tokens: token 0 and PT1 (M / 2: the second request's first row at spec c2) get their logits published first +NPT = _cfg("npt", 0) +PT1 = M_MAX // 2 +PF_PRIO = bool( + _cfg("pf_prio", 0) +) # L2 prefetch of the priority tokens' experts (this CTA's 1/NCTA of each) +PF_W1 = TWO_I * (H // 2) // NCTA +PF_S1 = TWO_I * (H // 16) // NCTA +PF_W2 = H * (I_TP // 2) // NCTA +PF_S2 = H * (I_TP // 16) // NCTA +assert PF_W1 % 16 == 0 and PF_S1 % 16 == 0 and PF_W2 % 16 == 0 and PF_S2 % 16 == 0 + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + # B rows n >= 32 (N1 = 64) go to the atom's next 32-bit column (row n: lane n % 32, column n // 32) + for w in cutlass.range_constexpr((N1 + 31) // 32): + n = lane + cutlass.Int32(32 * w) + q = n // cutlass.Int32(QS) + tl = n - q * cutlass.Int32(QS) + sf_ok = (n < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr( + base_row + warp * cutlass.Int32(32), + base_col + SFB1_COL + a * 4 + w * SFB1_WSTEP, + ), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + gpos, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = gpos % cutlass.Int32(SH_NST) # ring position (stage g of this pass) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue( + lane, + sid, + cutlass.Int32(g), + s_big, + sh_full, + tma_gu, + tma_dn, + gusf, + dnsf, + cutlass.Int32(g), + ) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + elif cutlass.const_expr(not SH_PF_POST): + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(SH_PF_POST): + if warp == 0: # after the grid dependency: DRAM is idle during the routed front + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + for ps in cutlass.range_constexpr(SH_NPASS): + TT0 = 0 # one pass: B rows SHN q + t (gate_up), t (down) + MS = M_MAX + GB = ps * SH_STAGES # ring positions before this pass + PAR = ps & 1 # parity of the per-pass barriers (B ready, accumulators) + if warp == 0: + for g in cutlass.range(SH_NST if ps == 0 else 0, SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) # ring position over the passes + _wait( + sh_empty.subview(gg % cutlass.Int32(SH_NST)), + ((gg // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf, gg) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=2 * SHN, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SHN, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) + st = gg % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (gg // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32((2 * SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32((SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(SHN), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + if cutlass.const_expr(ps > 0): + _wait( + acc_dn, PAR ^ 1 + ) # the previous pass's MMAs are done: its B tiles and TMEM scales are free + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((MS * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(MS * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=(t + cutlass.Int32(TT0)) * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(SHN) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + + it * cutlass.Int32(2 * SHN * 128) + + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of + # k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(SHN - 1) + ok = (n < cutlass.Int32(2 * SHN)) & (t < cutlass.Int32(MS)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n // cutlass.Int32(SHN)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t + # at column 8 q + t + _wait(acc_gu, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=2 * SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(MS): + bigf.store( + cutlass.Float32(av[SHN + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(MS): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((32 * MS + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(32 * MS): + p = pq // cutlass.Int32(MS) + tt = pq - p * cutlass.Int32(MS) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(SHN) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(32) + p) * cutlass.Int32(SHN) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview( + (tt + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + + sid * cutlass.Int32(32) + + p + ) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((MS * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(MS * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = (t + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + kb * cutlass.Int32( + 32 + ) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + + kt * cutlass.Int32(SHN * 128) + + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (t & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + ( + ((c0 + cutlass.Int32(1)) ^ (t & cutlass.Int32(7))) + << cutlass.Int32(2) + ), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(MS) + for kt in cutlass.range_constexpr(6): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * SHN, + cutlass.Float32, + ), + num=SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(MS): + _st_relaxed_u32( + dnx.subview(cutlass.Int32((tt + TT0) * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +# ============================================================================= +# v2 helpers: global release / acquire, proxy fences. +# ============================================================================= +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _atom_add_u64(addr, val, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr), cutlass.Int64(val)], + "atom.relaxed.gpu.global.add.u64 $0, [$1], $2;", + "=l,l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.proxy.async.global;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_acq_rel_gpu(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.acq_rel.gpu;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _dbg_delay(salt): + """(DBG_DELAY stress builds) a pseudo-random ~0..4 us stall: protocol robustness under skew.""" + if cutlass.const_expr(DBG_DELAY): + r = cutlass.Int32(_gtime() >> cutlass.Int64(6)) * cutlass.Int32(1103515245) + salt + if ((r >> cutlass.Int32(12)) & cutlass.Int32(7)) == cutlass.Int32(0): + for _ in cutlass.range((r >> cutlass.Int32(16)) & cutlass.Int32(31), unroll=1): + _nanosleep(128) + + +@cute.jit +def _waitt(bar: cutlass.Array, parity, acc): + """_wait that (TIMING) adds the waited ns to acc.""" + t0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + if cutlass.const_expr(TIMING): + acc = acc + (_gtime() - t0) + return acc + + +@cute.jit +def _tstore(tdbg: cutlass.Array, bx, i: cutlass.Constexpr, v): + if cutlass.const_expr(TIMING): + tdbg.store(cutlass.Int64(v), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@cute.jit +def _issue_fc1( + lane, + stage, + e, + c, + j, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a1, + tma_s1, + s1w: cutlass.Array, +): + """FC1 stage j of unit c of expert e into ring stage `stage` (the caller waited its empty barrier): 4 TMA boxes + (rows 32 c .., k-tiles 4 j + q) and the stage's scale atoms (SFA_CP: word [r][q] of atom kb <- raw atom + (k-tile 4 j + q, kb) word [r][c % 4], by this warp's 32 lanes).""" + row = c * cutlass.Int32(R1) + kx = j * cutlass.Int32(C1 * KT_BYTES) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview(stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(q * NUM_BYTES_BOX1)), + tma_a1, + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32(SF1_KA) + + j * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview(stage * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1, + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + atom0 = (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32( + SF1_KA + ) + j * cutlass.Int32(16) + src0 = atom0 * cutlass.Int32(128) + lane * cutlass.Int32(4) + (c & cutlass.Int32(3)) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _issue_fc2( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a2, + tma_s2, + s2w: cutlass.Array, +): + """FC2 k-tile t2 of output block blk for the group's n_here slots: box jj = rows 32 blk .. of slot jj's w2 at + k-tile t2; scale word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) word [r][blk % 4].""" + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32(NUM_BYTES_BOX1 if SFA_CP else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2, + (cutlass.Int32(t2 * KT_BYTES), blk * cutlass.Int32(R2), e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2, + ( + cutlass.Int32(0), + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _quant_words(w, sc, rsc, rsix): + """NVFP4 of one 16-value block given as 8 bf16-pair words: (lo word, hi word, scale e4m3 bits).""" + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.jit +def _quant_block(xs32: cutlass.Array, widx, sc, rsc, rsix): + """NVFP4 of one 16-value block (8 bf16 words at widx): (lo word, hi word, scale e4m3 bits), fp4_quantize's ops.""" + w = [] + for h in cutlass.range_constexpr(2): + v4 = xs32.load(idx=widx + cutlass.Int32(4 * h), vector_size=4, alignment=16) + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_gu_desc: cutlass.GridConstant[cuda.TensorMap], + tma_dn_desc: cutlass.GridConstant[cuda.TensorMap], + bias: cutlass.Array, # f32 [128] + escal: cutlass.Array, # f32 [6, 128] + plan: cutlass.Array, # int32 [4 M + 1, NCTA, M + 1] + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + img32: cutlass.Array, # int32 [NG_CAP, IMG_BYTES / 4]: per-group FC2 B image (tile + scale atoms) + cnt: cutlass.Array, # int32 [NG_CAP * 8]: per-group unit counters (index 8 g) + s1w: cutlass.Array, # int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # int32 words of w2_weight_scale + w1b: cutlass.Array, # int8 bytes of w3_w1_weight (L2 prefetch) + w2b: cutlass.Array, # int8 bytes of w2_weight (L2 prefetch) + epochs: cutlass.Array, # int32 [NCTA + NSH] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8] + tdbg: cutlass.Array, # int64 [NCTA + NSH, 64] + xb32: cutlass.Array, # int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # f32 router weight [128, 6144] + sfsc: cutlass.Array, # f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # int64 [LG_REPL, LG_STRIDE] + guw: cutlass.Array, + gusf: cutlass.Array, + dnw: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if cutlass.const_expr(TICKET): + # role by entry order: the CTAs that become resident last (behind the predecessor's CTAs) take the shared + # expert roles (most slack), the first NCTA the routed ones (the router GEMV / logits exchange) + s_tk = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + if tidx == 0: + old = _atom_add_u64( + cnt.subview(cutlass.Int32(NG_CAP * 8 + 8)).data_ptr().toint(), cutlass.Int64(1) + ) + s_tk.store(cutlass.Int32(old % cutlass.Int64(NCTA + NSH)), idx=0) + prims.barrier_cta_sync(0) + bx = s_tk.load(idx=0) + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + # ring scale words: SFA_CP: the instruction atoms [STAGES][4 x 512 B]; else raw boxes [STAGES][8 KB] + spliced + sSF = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(STAGES * STAGE_SF // 4,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA + (0 if SFA_CP else STAGES * NUM_BYTES_SFA_RAW)), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=512, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(NB2BUF * IMG_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), shape=(IMG_SF // 4,), dtype=cutlass.Int32, alignment=1024 + ) + ys = cutlass.Array( + cutlass.BFloat16, + NB2_CAP * M_MAX * TOP_K * R2, + space=cutlass.AddressSpace.smem, + alignment=16, + ) # bf16 values + s_part = cutlass.Array( + cutlass.Float32, SPB * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array( + cutlass.Float32, 9 * M_MAX + 16, space=cutlass.AddressSpace.smem, alignment=16 + ) + # the activation's scale codes [M, 384] until they are in TMEM: aliases ys (first written by the FC2 epilogue) + s_xsf = cutlass.Array( + ys.data_ptr(0), shape=(M_MAX * (H // 64),), dtype=cutlass.Int32, alignment=16 + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + img_bar = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + sfb2_free = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_full = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + x_empty = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + route_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + ys_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + grp_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_load = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + topk_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tok0_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + px_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + s_prog = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array(cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_gn = cutlass.Array( + cutlass.Int32, NG_CAP, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots per group + s_emap = cutlass.Array(cutlass.Int32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array(cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32) + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + if tidx < NB2BUF: + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), 1) + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < NXB: + prims.mbarrier_init(x_full.subview(tidx), 1) + prims.mbarrier_init(x_empty.subview(tidx), NWARPS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + prims.mbarrier_init(sfb2_free.subview(0), 1) + prims.mbarrier_init(route_bar.subview(0), 32) + prims.mbarrier_init(ys_done.subview(0), EPI_THREADS) + prims.mbarrier_init(b1_load.subview(0), 1) + prims.mbarrier_init(topk_bar.subview(0), M_MAX * TOP_K) + prims.mbarrier_init(tok0_bar.subview(0), TOP_K) + prims.mbarrier_init(px_full.subview(0), 1) + if tidx < 4: + prims.mbarrier_init(grp_done.subview(tidx), EPI_THREADS) + if tidx == 0: + s_prog.store(cutlass.Int32(0), idx=0) + if warp == X_WARP: # slot tables: empty (the slot warp fills them after the top-k) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + for z in cutlass.range_constexpr(E_LOCAL // 32): + s_emap.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + if cutlass.const_expr(not SFA_CP): + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + # router weights (expert bx) in registers: thread tg = tid % 256 owns elements 4 tg + 1024 j + tg = tidx % cutlass.Int32(GEMV_THREADS) + half = tidx // cutlass.Int32(GEMV_THREADS) + rwv = [ + rw.load( + idx=bx * cutlass.Int32(H) + tg * cutlass.Int32(4) + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + # ---- x chunks into the idle ring (warp 6): chunk c = K [XCK c, XCK (c + 1)) of every token + if warp == S_WARP: + if cutlass.const_expr(NPT > 0): + # the priority tokens' rows first (their logits are published ahead of the others) + if lane == 0: + prims.mbarrier_arrive_expect_tx(px_full.subview(0), NPT * H * 2) + cute.arch.sync_warp() + if lane < NPT: + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(lane * cutlass.Int32(H * 2)), + xb32.subview( + cutlass.select_( + lane == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1) + ) + * cutlass.Int32(H // 2) + ), + px_full.subview(0), + H * 2, + ) + for c in cutlass.range_constexpr(NXB): + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(c), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(c * XCB) + lane * cutlass.Int32(XCK * 2)), + xb32.subview(lane * cutlass.Int32(H // 2) + cutlass.Int32(c * XCK // 2)), + x_full.subview(c), + XCK * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + img8f = cutlass.Array( + img32.data_ptr(0), shape=(IMG_TOTAL,), dtype=cutlass.Int8, alignment=16 + ) + if cutlass.const_expr(NPT > 0): + # ---- priority tokens (0 and PT1): their logits first (m3_router_gemv's per-thread / warp order) + _wait(px_full.subview(0), 0) + ph = cutlass.select_(half < cutlass.Int32(NPT), half, cutlass.Int32(0)) + pacc = cutlass.Float32(0.0) + px32 = cutlass.Array( + sB1.data_ptr(0), shape=(NPT * H // 2,), dtype=cutlass.Int32, alignment=16 + ) + for j in cutlass.range_constexpr(H // 1024): + pxw = px32.load( + idx=ph * cutlass.Int32(H // 2) + cutlass.Int32(512 * j) + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + px0 = cutlass.Int32(pxw[0]) + px1 = cutlass.Int32(pxw[1]) + pacc = _fma_rn( + (px0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][0]), + pacc, + ) + pacc = _fma_rn( + (px0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][1]), + pacc, + ) + pacc = _fma_rn( + (px1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][2]), + pacc, + ) + pacc = _fma_rn( + (px1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][3]), + pacc, + ) + for off in (16, 8, 4, 2, 1): + pacc = _add_rn(pacc, cute.arch.shuffle_sync_bfly(pacc, off)) + if lane == 0: + s_red.store(pacc, idx=cutlass.Int32(9 * M_MAX) + warp) # [half][8 warps] + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if tidx < cutlass.Int32(NPT * LG_REPL): + hh = tidx // cutlass.Int32(LG_REPL) + rr = tidx - hh * cutlass.Int32(LG_REPL) + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn( + sacc, s_red.load(idx=cutlass.Int32(9 * M_MAX + w) + hh * cutlass.Int32(8)) + ) + tp = cutlass.select_(hh == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1)) + _st_relaxed_u64( + lgbuf.subview(rr * cutlass.Int32(LG_STRIDE) + tp * cutlass.Int32(E_LOCAL) + bx) + .data_ptr() + .toint(), + ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)), + ) + if tidx == 0: + _stamp(tdbg, bx, 43) + # ---- this CTA's share of the NVFP4 activation (QB blocks, from global) for the global B1 image: loads first + QB = (M_MAX * (H // 16) + NCTA - 1) // NCTA + qb = bx * cutlass.Int32(QB) + tidx + q_ok = (tidx < cutlass.Int32(QB)) & (qb < cutlass.Int32(M_MAX * (H // 16))) + qb = cutlass.select_(q_ok, qb, cutlass.Int32(0)) + qtok = qb // cutlass.Int32(H // 16) + qkb = qb - qtok * cutlass.Int32(H // 16) + xq_v = [ + xb32.load( + idx=qtok * cutlass.Int32(H // 2) + qkb * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + for h in range(2) + ] + # ---- router GEMV from the staged chunks, all 512 threads (rolled: the front runs once per call) + MH = (M_MAX + 1) // 2 # tokens per GEMV half + acc_r = cute.make_rmem_tensor((MH,), cutlass.Float32) + for i in cutlass.range_constexpr(MH): + acc_r[i] = cutlass.Float32(0.0) + big32 = cutlass.Array( + s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16 + ) + JPC = XCK // 1024 # router j-steps per chunk + for xc in cutlass.range(NXC, unroll=1): + buf = xc % cutlass.Int32(NXB) + _wait(x_full.subview(buf), (xc // cutlass.Int32(NXB)) & cutlass.Int32(1)) + if (tidx == 0) & (xc == cutlass.Int32(0)): + _stamp(tdbg, bx, 41) + xw0 = buf * cutlass.Int32(XCB // 4) # chunk word base in s_big + for jj in cutlass.range_constexpr(JPC): + jg = xc * cutlass.Int32(JPC) + cutlass.Int32(jj) + wsel = [cutlass.Float32(rwv[0][z]) for z in range(4)] + for jx in cutlass.range_constexpr(1, H // 1024): + for z in cutlass.range_constexpr(4): + wsel[z] = cutlass.select_( + jg == cutlass.Int32(jx), cutlass.Float32(rwv[jx][z]), wsel[z] + ) + for i in cutlass.range_constexpr(MH): + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + xw = big32.load( + idx=xw0 + + t * cutlass.Int32(XCK // 2) + + cutlass.Int32(512 * jj) + + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + acc_g = acc_r[i] + acc_g = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[0], acc_g + ) + acc_g = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[1], acc_g + ) + acc_g = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[2], acc_g + ) + acc_g = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[3], acc_g + ) + acc_r[i] = acc_g + if cutlass.const_expr(NXB < NXC): + if xc + cutlass.Int32(NXB) < cutlass.Int32(NXC): + # every thread's reads of this chunk are done (CTA barrier, then a generic -> async proxy fence by + # the refilling warp: the bulk copy overwrites what the generic proxy read); racecheck-clean + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == S_WARP: + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(buf), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + buf * cutlass.Int32(XCB) + lane * cutlass.Int32(XCK * 2) + ), + xb32.subview( + lane * cutlass.Int32(H // 2) + + (xc + cutlass.Int32(NXB)) * cutlass.Int32(XCK // 2) + ), + x_full.subview(buf), + XCK * 2, + ) + if tidx == 0: + _stamp(tdbg, bx, 42) + for i in cutlass.range_constexpr(MH): + v = acc_r[i] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + if (lane == 0) & (t < cutlass.Int32(M_MAX)): + s_red.store(v, idx=(warp % cutlass.Int32(8)) * cutlass.Int32(M_MAX) + t) + # this CTA's quantization blocks into the global B1 image (B1 tile bytes, then the [M, 384] scale codes) + if q_ok: + qw = [] + for h in cutlass.range_constexpr(2): + for z in cutlass.range_constexpr(4): + qw.append(cutlass.Int32(xq_v[h][z])) + qlo, qhi, qsf = _quant_words(qw, sc, rsc, rsix) + img32.store( + (qlo, qhi), idx=cutlass.Int32(B1IMG_OFF // 4) + _b1_addr(qtok, qkb), alignment=8 + ) + img8f.store(cutlass.Int8(qsf), idx=cutlass.Int32(B1IMG_OFF + S1 * NUM_BYTES_B1) + qb) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == cutlass.Int32(0): + _dbg_delay(bx * cutlass.Int32(3)) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX + t))) + lw = ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)) + _st_relaxed_u64( + lgbuf.subview(tidx * cutlass.Int32(LG_STRIDE) + cutlass.Int32(t * E_LOCAL) + bx) + .data_ptr() + .toint(), + lw, + ) + if tidx == 0: + _stamp(tdbg, bx, 21) + if tidx == cutlass.Int32( + THREADS - 32 + ): # B1 share released after the logits (its gpu fence stalls warp 15 only) + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() # once, after the join (ALL_BAR) of the writer threads + _red_release_add( + cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint(), cutlass.Int32(1) + ) + # ---- the B1 image (every CTA's share) into smem (warp 13; the top-k warps are 0 .. M - 1) + if cutlass.const_expr(B1_WARP >= 0): + if warp == cutlass.Int32(B1_WARP): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + + # ---- top-k: token t on warp t (keys sigmoid + bias, ties to the lower id), TRTLLM-Gen MiniMax2 ops + tk_t = cutlass.Int32(-1) # this warp's top-k token (TK_WARPS[t] = warp) + for i_tk in cutlass.range_constexpr(M_MAX): + tk_t = cutlass.select_(warp == cutlass.Int32(TK_WARPS[i_tk]), cutlass.Int32(i_tk), tk_t) + if tk_t >= cutlass.Int32(0): + t = tk_t + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 22) + lgv = [cutlass.Int32(res[z]).bitcast(cutlass.Float32) for z in range(4)] + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hk = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hk[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + if lane < TOP_K: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 20) + if cutlass.const_expr(PF_PRIO): + if (t == cutlass.Int32(0)) | (t == cutlass.Int32(PT1)): + if lane < TOP_K: + # this CTA's 1/NCTA of each of the token's experts into L2 (DRAM is idle until the routing is + # complete; every byte is streamed later anyway) + _l2_prefetch( + w1b.subview( + e_l * cutlass.Int32(TWO_I * H // 2) + bx * cutlass.Int32(PF_W1) + ) + .data_ptr() + .toint(), + PF_W1, + ) + _l2_prefetch( + s1w.subview( + (e_l * cutlass.Int32(TWO_I * H // 16) + bx * cutlass.Int32(PF_S1)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S1, + ) + _l2_prefetch( + w2b.subview( + e_l * cutlass.Int32(H * I_TP // 2) + bx * cutlass.Int32(PF_W2) + ) + .data_ptr() + .toint(), + PF_W2, + ) + _l2_prefetch( + s2w.subview( + (e_l * cutlass.Int32(H * I_TP // 16) + bx * cutlass.Int32(PF_S2)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S2, + ) + cute.arch.sync_warp() + if lane < TOP_K: # every writer lane arrives (its own release of s_ids / s_w) + prims.mbarrier_arrive(topk_bar.subview(0)) # this token's ids / weights are in smem + if t == cutlass.Int32(0): + prims.mbarrier_arrive( + tok0_bar.subview(0) + ) # token 0's experts: the early group-0 stages + cute.arch.sync_warp() + + if cutlass.const_expr(EARLY0): + if warp == W_WARP: + # slots 0..3 = token 0's four distinct experts; group 0's stage ranges are the same for every G + _wait(tok0_bar.subview(0), 0) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(STAGES): + _issue_fc1( + lane, + n0, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + n0 = n0 + cutlass.Int32(1) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if warp == X_WARP: + # slots: entry i < 4 M is (token i / 4, top-k i % 4); distinct experts in entry order + NCH = (TOP_K * M_MAX + 31) // 32 + _wait(topk_bar.subview(0), 0) + gcnt = cutlass.Int32(0) + rec = [] # per chunk: (first, rank, slot rank of the entry, expert, entry valid) + for ch in cutlass.range_constexpr(NCH): + ei = lane + cutlass.Int32(32 * ch) + valid = ei < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, ei, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + known = cutlass.Boolean(False) + prev = cutlass.Int32(-1) + key = my_e + if cutlass.const_expr(ch > 0): + prev = cutlass.select_( + valid, + s_emap.load(idx=cutlass.select_(valid, my_e, cutlass.Int32(0))), + cutlass.Int32(-1), + ) + known = prev >= cutlass.Int32(0) + key = cutlass.select_(known, cutlass.Int32(-1000) - lane, my_e) + cute.arch.sync_warp() # every lane's s_emap read before any lane's write + same = _match_any(key) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (known == cutlass.Boolean(False)) & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = gcnt + cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + slot_l = cutlass.select_(known, prev, slot_l) + if first: + s_emap.store(rank, idx=my_e) + if valid: + if bx == 0: + route_dbg.store( + my_e, + idx=(ei // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + ei % cutlass.Int32(TOP_K), + ) + rec.append((first, rank, slot_l, my_e, valid)) + cute.arch.sync_warp() + gcnt = gcnt + cute.arch.popc(bal) + # slot of rank r: with a partial last group (tail < 4) and >= 3 groups, that group becomes group 1 (its + # FC2 then is not the last hand-off behind a short FC1 phase); the full groups 1.. shift up by one + ngr = (gcnt + cutlass.Int32(3)) >> cutlass.Int32(2) + tail = gcnt - (ngr - cutlass.Int32(1)) * cutlass.Int32(4) + do_perm = ( + cutlass.Boolean(PERM) + & (tail < cutlass.Int32(4)) + & (ngr >= cutlass.Int32(PERM_MIN_NG)) + ) + for z in cutlass.range_constexpr((G_CAP + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP): + s_el.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + for ch in cutlass.range_constexpr(NCH): + first, rank, slot_r, my_e, valid = rec[ch] + ei = lane + cutlass.Int32(32 * ch) + for which in cutlass.range_constexpr(2): + r = rank if which == 0 else slot_r + sl_m = cutlass.select_( + r < cutlass.Int32(4), + r, + cutlass.select_( + r >= (ngr - cutlass.Int32(1)) * cutlass.Int32(4), + r - (ngr - cutlass.Int32(2)) * cutlass.Int32(4), + r + cutlass.Int32(4), + ), + ) + sl_m = cutlass.select_(do_perm, sl_m, r) + if cutlass.const_expr(which == 0): + if first: + s_el.store(my_e, idx=sl_m) + else: + if valid: + s_kpos.store( + ei % cutlass.Int32(TOP_K), + idx=sl_m * cutlass.Int32(M_MAX) + ei // cutlass.Int32(TOP_K), + ) + # slots per group + if lane < cutlass.Int32(NG_CAP): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + nh = cutlass.select_(do_perm & (lane == cutlass.Int32(1)), tail, nh) + nh = cutlass.select_( + do_perm & (lane == ngr - cutlass.Int32(1)), cutlass.Int32(4), nh + ) + s_gn.store(nh, idx=lane) + cute.arch.sync_warp() + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + prims.mbarrier_arrive( + route_bar.subview(0) + ) # every lane: its slots / kpos / plan-row writes are in smem + cute.arch.sync_warp() + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = s_gn.load(idx=lane) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add(cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad) + else: + _wait(route_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(B1_WARP < 0): + if warp == cutlass.Int32(0): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + if warp < 4: # B1 scales to TMEM (the MMA needs them only with the first ring stage) + _wait(b1_load.subview(0), 0) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + n_iss = cutlass.Int32( + 0 + ) # stages issued (early ones included): the L2-prefetch warp's throttle + w_emp = cutlass.Int64(0) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _dbg_delay(bx) + _issue_fc1( + lane, + stage, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _issue_fc2( + lane, + stage, + g, + blk, + t2, + n_here, + s_el, + sA, + sSF, + ab_full, + tma_a2_desc.get_ptr(), + tma_s2_desc.get_ptr(), + s2w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if lane == 0: + _stamp(tdbg, bx, 3) + _tstore(tdbg, bx, 30, w_emp) + + # ---- (SFA_CP = 0) scale splicer (warp 6): atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if cutlass.const_expr(not SFA_CP): + if warp == S_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt_ = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt_ = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt_ = cutlass.Int32(KT2) + for _j in cutlass.range(cnt_, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32( + 4 + ) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSF.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSF.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSF.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSF.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + # ---- FC2 B loader (warp 5): group g's image once its counter is complete + if warp == X_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for g in range(ng): + buf = g % cutlass.Int32(NB2BUF) + if g >= cutlass.Int32(NB2BUF): + _waitd( + b2_free.subview(buf), + ((g // cutlass.Int32(NB2BUF)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 29) + caddr = cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2_POLL_NS) + cv = _ld_acquire_u32(caddr) + cute.arch.sync_warp() + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 51) + _fence_proxy_async_global() + if g >= cutlass.Int32(1): + _waitd( + sfb2_free.subview(0), (g - cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 107 + ) + if lane == 0: + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(0), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), tdbg, 108 + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 52) + if g >= cutlass.Int32(1): + # the single TMEM scale slot: the previous group's FC2 MMAs are done + gp = g - cutlass.Int32(1) + _waitd( + b2_free.subview(gp % cutlass.Int32(NB2BUF)), + (gp // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), + tdbg, + 109, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 53) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for aa in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFB2_COL + aa * 4), + sd0 + cutlass.Int64(aa * (512 >> 4)), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + prims.tcgen05_commit(sfb2_free.subview(0)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + + # ---- (PFD) L2 prefetch warp (12): the same stage sequence as the producer, PFD stages ahead of it + if cutlass.const_expr(PFD > 0): + if warp == cutlass.Int32(PF_WARP): + pf = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + kx = j * cutlass.Int32(C1 * KT_BYTES) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_prefetch( + tma_a1_desc.get_ptr(), + [ + kx + cutlass.Int32(q * KT_BYTES), + c * cutlass.Int32(R1), + e, + ], + [], + ) + atom0 = ( + e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2)) + ) * cutlass.Int32(SF1_KA) + j * cutlass.Int32(16) + _l2_prefetch( + s1w.subview(atom0 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 8192, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load( + idx=g * cutlass.Int32(4) + cutlass.Int32(jj) + ) + prims.cp_async_bulk_tensor_prefetch( + tma_a2_desc.get_ptr(), + [ + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ], + [], + ) + a2 = ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) * cutlass.Int32(SF2_KA) + cutlass.Int32(4 * t2) + _l2_prefetch( + s2w.subview(a2 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 2048, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + + # ---- (REL_WARP) group releases (warp 8): this CTA's finished units of group p, one release add per phase + if cutlass.const_expr(REL_WARP): + if warp == cutlass.Int32(8): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + grp_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + _dbg_delay(bx * cutlass.Int32(13) + p) + if (lane == 0) & (n_done > cutlass.Int32(0)): + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + _red_release_add( + cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + if (lane == 0) & (p == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 54) + cute.arch.sync_warp() + + # ---- MMA (warp 7) + if warp == MMA_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + m_b2 = cutlass.Int64(0) + m_full = cutlass.Int64(0) + m_acc = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + m_acc = _waitt( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + cute.arch.sync_warp() + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g % cutlass.Int32(NB2BUF) + m_b2_0 = m_b2 + m_b2 = _waitt( + b2_ready.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), m_b2 + ) + if cutlass.const_expr(TIMING): + if (lane == 0) & (g < cutlass.Int32(4)): + tdbg.store( + m_b2 - m_b2_0, + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(44) + g, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + m_acc = _waitt( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB2_COL + cutlass.Int32(t2 * NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + cute.arch.sync_warp() + k2 = k2 + cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit( + b2_free.subview(buf) + ) # this group's B2 tile and TMEM scales are free + cute.arch.sync_warp() + if lane == 0: + _tstore(tdbg, bx, 26, m_b2) + _tstore(tdbg, bx, 27, m_full) + _tstore(tdbg, bx, 28, m_acc) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + img8 = cutlass.Array( + img32.data_ptr(0), shape=(NG_CAP * IMG_BYTES,), dtype=cutlass.Int8, alignment=16 + ) + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if qi < cutlass.Int32(6 * G_CAP): + e_sl = s_el.load(idx=sl_e) + if e_sl >= cutlass.Int32(0): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32(escal.load(idx=fld * cutlass.Int32(E_LOCAL) + e_sl)), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + e_wait = cutlass.Int64(0) + e_busy = cutlass.Int64(0) + e_a = cutlass.Int64(0) + e_pw = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + unit_done = (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)) + e_wait = _waitt( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + tb0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1_LD, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=(t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + tc0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_a = e_a + (tc0 - tb0) + g_u = slot >> cutlass.Int32(2) + jrow = slot & cutlass.Int32(3) + # this warp's tokens t = warp + 4 tr, all in flight together (latency-bound otherwise) + NT = (M_MAX + 3) // 4 + xr = cute.make_rmem_tensor((NT,), cutlass.Float32) + for tr in cutlass.range_constexpr(NT): + tq = cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + xr[tr] = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=(tq * C1 + 0) * R1 + lane), + s_part.load(idx=(tq * C1 + 1) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 2) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 3) * R1 + lane), + ) + if (j0 == cutlass.Int32(0)) & (j1 != cutlass.Int32(S1)): + # head of a unit split with the next CTA: its partial sums + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + _st_relaxed_u32( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint(), + xr[tr].bitcast(cutlass.Int32), + ) + else: + if j0 != cutlass.Int32(0): + # tail: the head's partial sums (head first, as TRTLLM-Gen's K order) + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + pbits = _ld_relaxed_u32(paddr) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xr[tr] = _add_rn(pbits.bitcast(cutlass.Float32), xr[tr]) + cute.arch.sync_warp() + _dbg_delay(bx * cutlass.Int32(7) + warp) + es = slot * cutlass.Int32(8) + ea = s_esc.load(idx=es) + ec = s_esc.load(idx=es + cutlass.Int32(1)) + eal = s_esc.load(idx=es + cutlass.Int32(2)) + ebe = s_esc.load(idx=es + cutlass.Int32(3)) + eli = s_esc.load(idx=es + cutlass.Int32(4)) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + n0 = jrow * cutlass.Int32(M_MAX) + t2 = cblk >> cutlass.Int32(4) + cc = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = g_u * cutlass.Int32(IMG_BYTES) + sfoff = ( + gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2 * cutlass.Int32(NUM_KB) + + ((cblk & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(512) + ) + + (cblk & cutlass.Int32(3)) + ) + for tr in cutlass.range_constexpr(NT): + # (skip) only tokens routed to this slot (warp-uniform); other B2 rows are never used + if s_kpos.load( + idx=slot * cutlass.Int32(M_MAX) + + cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + ) >= cutlass.Int32(0): + xv = xr[tr] + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + qv, sfb = _swiglu_quant( + xv, gate_raw, is_up, ea, ec, eal, ebe, eli + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + t = warp + cutlass.Int32(4 * tr) + if ( + (lane == 0) + & (t < cutlass.Int32(M_MAX)) + & cutlass.Boolean(not DBG_NOSTORE) + ): + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), row n + # at byte 16 (n % 32) + 4 (n / 32), byte c % 4 + n = n0 + t + img32.store( + (lo, hi), + idx=( + gbase + + t2 * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=sfoff + + (n & cutlass.Int32(31)) * cutlass.Int32(16) + + (n >> cutlass.Int32(5)) * cutlass.Int32(4), + ) + tc1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_pw = e_pw + (tc1 - tc0) + if cutlass.const_expr(not REL_WARP): + if unit_done: + if lane == 0: + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + if cutlass.const_expr(EPI_FENCE): + _fence_acq_rel_gpu() + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=EPI_BAR, number_of_threads=EPI_THREADS + ) # s_part reads done + if cutlass.const_expr(not REL_WARP): + if unit_done & (tidx == 0): + _red_release_add( + cnt.subview(g_u * cutlass.Int32(8)).data_ptr().toint(), + cutlass.Int32(1), + ) + if cutlass.const_expr(TIMING): + e_busy = e_busy + (_gtime() - tb0) + k1 = k1 + cutlass.Int32(1) + if cutlass.const_expr(REL_WARP): + prims.mbarrier_arrive( + grp_done.subview(p & cutlass.Int32(3)) + ) # this CTA's group-p units are out + if tidx == 0: + _stamp(tdbg, bx, 10) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + e_wait = _waitt( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + if cutlass.const_expr(TIMING): + if ( + (tidx == 0) + & (g == ng - cutlass.Int32(1)) + & (bb == cutlass.Int32(0)) + ): + _stamp(tdbg, bx, 8) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = [] + ya = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), + cutlass.Float32, + ), + num=ACC2_STRIDE, + ) + for t in cutlass.range_constexpr(M_MAX): + v = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + v = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(ya[w * M_MAX + t]), v + ) + yv.append(v) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if warp < s_gn.load(idx=g): + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + cutlass.BFloat16( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)) + ), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + if tidx == 0: + _stamp(tdbg, bx, 6) + _tstore(tdbg, bx, 31, e_wait) + _tstore(tdbg, bx, 34, e_busy) + _tstore(tdbg, bx, 36, e_a) + _tstore(tdbg, bx, 38, e_pw) + prims.mbarrier_arrive(ys_done.subview(0)) # this thread's ys values are in smem + + # ---- finalize on all 16 warps: (block bb, token t) pairs, the fma chain over t's top-4 in order, + y_sh + NPAIR = NB2_CAP * M_MAX + NPW = (NPAIR + NWARPS - 1) // NWARPS + wyl = [] + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + ok = (pair < cutlass.Int32(NPAIR)) & (bb < nb) + ch = cutlass.select_( + ok, + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wyl.append( + _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + if cutlass.const_expr(SHARED) + else cutlass.Int32(0) + ) + _wait(ys_done.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 19) + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if (pair < cutlass.Int32(NPAIR)) & (bb < nb): + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ), + acc_o, + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SHARED): + wy = wyl[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + bias: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + img32: cute.Tensor, + cnt: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + w1b: cute.Tensor, + w2b: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + bias, + escal, + plan, + partials, + img32, + cnt, + s1w, + s2w, + w1b, + w2b, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py new file mode 100644 index 000000000000..ac9892f31ed2 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py @@ -0,0 +1,3714 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP rank: ``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 +I_TP = 768 +TWO_I = 2 * I_TP +E_LOCAL = 128 +TOP_K = 4 +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +SH_I = 768 +SH_KB = H // 32 # gate_up K blocks +SH_DKB = SH_I // 32 # down K blocks +SH_SPLIT = 2 # gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges): MMA rows 64 q + .. = slice q +SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * (SH_I // 128) # + 2 down tiles x 6 k-tiles +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its 36 stages, each the 16 KB +# shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction scale +# atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 else 0) +) # (in-chain: a win at M <= 2, not at 4..8) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # 6 stages per unit +NUM_BYTES_A = MMA_M * KT_BYTES +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = I_TP // 256 # 3 +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +PARTIAL_EMPTY = -1 +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 +X_WARP = 5 +S_WARP = 6 +MMA_WARP = 7 +E0_WARP = MMA_WARP # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and NRW0 <= 3 and bool(_cfg("e0_off", 1)) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +assert M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +_SH_FIX = ( + SH_GT * 16 * 128 + + 6 * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + 64 * 8 * 4 + + 2 * 32 * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * 16 * 128 +_SB_XC = _SB_BH + 6 * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + 64 * 8 * 4 +_SB_END = _SB_GU + 2 * 32 * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N 16), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = 16 +SH_SFA = 32 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +SH_TMEM = 256 +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=16, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(2048 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(2048) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(16)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((32 * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(32 * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(32) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(32) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(6): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview((rec0 + qi) * cutlass.Int32(4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py new file mode 100644 index 000000000000..12f6008d439f --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py @@ -0,0 +1,602 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP4 rank, over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +HIDDEN = 6144 +INTER = 768 +NUM_EXPERTS = 128 +SHARED_INTER = 768 # the shared expert's intermediate per TP rank +SHARED_CTAS = 24 # m3_moe's shared-expert CTAs (beside num_ctas routed CTAs) +TOP_K = 4 +U1 = 2 * INTER // 32 # FC1 units per expert +MIN_CTAS = 64 # the router GEMV takes at most 2 experts per CTA (router weights held in registers) + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # FC1 stages (items) per unit +_ITEM_BYTES = _UNIT_BYTES // _S1 +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters (7 and 8 tokens). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 1, 1) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit). 6 at 8 tokens measured 0.4-0.5 us faster than the +# default 7 (same MMAs, bit-identical outputs). +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 6) +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). Stages 0..23: gate_up k-tiles + g (slice 0) / 24 + g (slice 1) of rows [gate 32 s .., up 768 + 32 s .., gate (slice 1), up (slice 1)]; 24..35: + down tile (g - 24) / 6 (rows 256 s + 128 j ..), k-tile (g - 24) % 6. Built once per layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty(n_cta, 36, 16384 + 512, dtype=torch.uint8, device=dev) + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = torch.where((q & 1) == 0, 32 * s + ri, I_ + 32 * s + ri) + sl = q >> 1 + for g in range(24): + kt = sl * 24 + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(12): + jt, kt = gd // 6, gd % 6 + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, 24 + gd, :16384] = img.reshape(-1) + out[s, 24 + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def make_phase_plan( + num_ctas: int, m_max: int, lag: int = 2, tail_w: float = 1.0, early0: bool = False +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work. Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [max(_S1, (limit - f) // _ITEM_BYTES) for f in fc2] + + lo, hi = 0, max(fc2) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + ): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, num_tokens, M3_MOE_LAG, 0.0 if num_tokens == 1 else 1.0, True + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + } + if shared: + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full((mod.G_CAP * U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw) + self.hbuf32 = torch.zeros(mod.G_CAP * num_tokens * U1 * 4, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [1536, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 768], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + key = (_ENGINE_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState` (or op16's state): ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py new file mode 100644 index 000000000000..446776419315 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host side of ``m3_moe16`` (m3_moe16_kernel.py): the MiniMax-M3 MoE layer (fused router + routed experts + shared +expert) for 9..16 decode tokens on one TP4 rank. The 1..8-token engine (op.M3MoeState) is untouched; select this +state only for num_tokens > 8. + +Output: bf16 [M, 6144] = bf16(shared + routed), the routed part bit for bit TRTLLM-Gen's fused MoE on the fp32 router +GEMV logits (m3_router_gemv's math), the shared part as the in-model m3_mx_gemv shared expert. + + state = M3Moe16State(device, num_tokens=12) # before CUDA-graph capture; one per token count + layer = state.layer(w3_w1_weight, w3_w1_weight_scale, w2_weight, w2_weight_scale, fc31_alpha, fc31_scale_c, + act_alpha, act_beta, act_clamp, fc2_alpha, e_score_correction_bias, + router_weight=router_weight, fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf)) + out = layer(hidden) # first call compiles: call once before capture + +Protocol state (per state object, shared by the layers that run on it in stream order): per-group FC2-B images and +unit counters, the B1 image + counter, per-CTA epochs. Every call adds exactly 192 to each group counter and NCTA to +the B1 counter, so nothing is reset between calls (graph replay safe). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch + +from .op import ( + HIDDEN, + SHARED_CTAS, + SHARED_INTER, + U1, + M3MoeLayer, + _compiled, + _kernel_module, + _shared_views, + _view, + is_supported, +) + + +def _group_slots(g: int, p: int, perm: bool, min_ng: int = 5) -> int: + """Slots of group p when the call has g distinct experts (the kernel's slot order).""" + ng = (g + 3) // 4 + tail = g - 4 * (ng - 1) + if p >= ng: + return 0 + if perm and tail < 4 and ng >= max(3, min_ng): + return tail if p == 1 else 4 + return min(4, g - 4 * p) + + +def make_phase_plan16(num_ctas: int, m_max: int, lag: int = 2, perm: bool = True, min_ng: int = 5): + """op.make_phase_plan (early0, tail_w 1) with the kernel's group sizes (PERM: a partial last group is group 1).""" + S1 = 6 + UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) + ITEM = UNIT_BYTES // S1 + blk_bytes = 32 * (768 // 2 + 768 // 16) + gcap = 4 * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + for g in range(4, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = _group_slots(g, p, perm, min_ng) + total = n_p * U1 * S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * _group_slots(g, p - lag, perm, min_ng) * blk_bytes + if p == ng - 1 and not (p == 0 and m_max >= 2): + for q in range(max(0, ng - lag), ng): + fc2[c] += nbs[c] * _group_slots(g, q, perm, min_ng) * blk_bytes + items = [0] * num_ctas + if total >= S1 * num_ctas: + + def counts_for(limit): + return [max(S1, (limit - f) // ITEM) for f in fc2] + + lo, hi = 0, max(fc2) + total * ITEM + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > S1), + key=lambda c: (-(fc2[c] + items[c] * ITEM), c), + ) + items[c] -= 1 + extra -= 1 + else: + for c in sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // S1]): + items[c] = S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +_ENGINE16_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe16_kernel.py") +M16_MIN_TOKENS, M16_MAX_TOKENS = 9, 16 +# Default routed CTAs (beside the SHARED_CTAS shared-expert CTAs). +M16_NUM_CTAS = 128 +# Phase lag: phase p streams FC1 of group p and FC2 of group p - M16_LAG (the kernel's default lag). +M16_LAG = 3 + + +class M3Moe16State: + """Workspace + compiled kernel of ``m3_moe16`` for ``num_tokens`` (9..16) tokens per call.""" + + fused_router = True + shared = True + + def __init__(self, device: torch.device, num_tokens: int, num_ctas: Optional[int] = None): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3Moe16State allocates and compiles: build it before CUDA-graph capture" + ) + if not M16_MIN_TOKENS <= num_tokens <= M16_MAX_TOKENS: + raise ValueError( + f"m3_moe16 takes {M16_MIN_TOKENS}..{M16_MAX_TOKENS} tokens per call, not {num_tokens}" + ) + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe16: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or M16_NUM_CTAS + if not 128 <= self.num_ctas <= sms - SHARED_CTAS: + raise ValueError( + f"m3_moe16 runs 128 .. {sms - SHARED_CTAS} routed CTAs + {SHARED_CTAS} shared (all resident)" + ) + self.m = num_tokens + plan, nb_cap = make_phase_plan16(self.num_ctas, num_tokens, M16_LAG) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "shared": 1, + } + try: + self.mod = mod = _kernel_module(self.config, _ENGINE16_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe16: num_tokens={num_tokens} does not fit the smem / TMEM budget: {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + # split units' head sums: every slot empty (0xFFFFFFFF) between calls + self.partials = torch.full((mod.G_CAP * U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw) + # FC2-B images per group + the B1 image; counters: group g at 8 g, B1 at 8 NG_CAP, CTA tickets at 8 NG_CAP + 8 + self.img = torch.zeros(mod.IMG_TOTAL // 4, dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 16, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros( + mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw + ) # logits: tag << 32 | f32 + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> M3MoeLayer: + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + layer.shared_weights = _shared_views(shared_weights, "m3_moe16") + layer.shared_mxfp8 = tuple(shared_weights) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError(f"m3_moe16: expects contiguous bf16 hidden states [{m}, {HIDDEN}]") + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + layer.bias, + layer.escal, + self.plan.view(-1), + self.partials, + self.img, + self.cnt, + layer.w3_w1_weight_scale.reshape(-1).view(torch.int32), + layer.w2_weight_scale.reshape(-1).view(torch.int32), + layer.w3_w1_weight.reshape(-1).view(torch.int8), + layer.w2_weight.reshape(-1).view(torch.int8), + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + *layer.shared_weights, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_moe16 compiles on its first call: call it once before capture" + ) + import cutlass.cute as cute + + sig = [_view(t, 16 if t.element_size() * t.numel() % 16 == 0 else 4) for t in args] + key = (_ENGINE16_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py new file mode 100644 index 000000000000..ef3a4c5847ea --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) as a CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..4247a736e8ca --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows per CTA = the MMA's M +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row +MMA_INST_K = 32 +K_BLOCKS = CTA_K // MMA_INST_K # 4 MMAs per k-tile, one per scale byte of the atom +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +# The swigluoai activation of MiniMax-M3's dense MLP, gate * sigmoid(1.702 * gate) * (up + 1) with gate clamped to <= 7 +# and up to [-7, 7], in exactly the instruction sequence the Triton silu_and_mul kernel compiles to (modules/swiglu.py +# on SM107): the gate clamp on the packed bf16 pair, the up clamp as min.xorsign.abs, two rounded multiplies into +# ex2.approx, div.full for the reciprocal, two rounded multiplies, round to bf16. No multiply feeds an add, so the +# explicit .rn forms compile to what the Triton kernel's contractable forms do. +SWIGLU_OAI_ALPHA = 1.702 +SWIGLU_OAI_LIMIT = 7.0 + + +@dsl_user_op +def _swiglu_oai_bf16x2(gate_word, up_word, *, loc=None, ip=None): + """The activation of one packed bf16 pair of gate and up values: its two bf16 results as one int32 word (the low + half from the low halves).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [cutlass.Int32(w).ir_value(loc=loc, ip=ip) for w in (gate_word, up_word)], + "{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, $1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, $2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 $0, r1, r0;\n\t}", "=r,r,r", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[ + bool + ], # x_bf16: x_words is [M, 2K] gate | up; the activation is swigluoai(gate, up) +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in if x_swiglu else k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + if cutlass.const_expr(x_swiglu): + # The up half of the same columns, K values after the gate in the row. + qz_wvecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(k_in // 2 + qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + if cutlass.const_expr(x_swiglu): + qz_word = _swiglu_oai_bf16x2( + qz_word, cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + ) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[bool] = False, +) -> None: + # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. + assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( + "x_swiglu needs x_bf16 and excludes din" + ) + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, x_swiglu, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py new file mode 100644 index 000000000000..041dc12b1ffd --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py @@ -0,0 +1,521 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) for already-quantized activations of at most 8 tokens. + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring): GB300 exposed time behind a 4.25 us PDL predecessor at M 1 / 8, vs the fused low-M GEMM: +# QKV 2.78 / 3.16 vs 4.58 / 6.21 us, o_proj 2.23 / 2.41 vs 4.40 / 5.09, shared gate_up 2.43 / 2.74 vs 3.56 / 4.30, +# shared down 1.74 / 1.66 vs 2.92 / 3.38. A config is (split, ring[, contig, nacc, a_tmem, st_async]); the optional +# fields default to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + (2560, 6144): (4, 12), + (6144, 2048): (2, 8), + # Shared-expert gate_up: split 2, ring 12, single accumulator, A operand from shared memory. + (1536, 6144): (2, 12, 1, 1, 0, 1), + # Shared-expert down: unsplit, ring 3, single accumulator, A operand from shared memory. + (6144, 768): (1, 3, 1, 1, 0, 1), + (2304, 6144): (4, 12), + (6144, 3072): (2, 12), + # Dense-layer gate_up: 6.4 / 6.3 us at M 1 / 8 vs FlashInfer 8.8. + (6144, 6144): (2, 12), +} +# Per-token-count configs (N, K, M): the grid that suits the concurrent routed-expert engine depends on how long +# that engine runs. At 2 tokens the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = { + (1536, 6144, 2): (4, 12, 1, 1, 0, 1), +} +# Above 8 tokens (spec-decode verify rows): the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = { + (1536, 6144): (4, 12, 1, 1, 0, 1), +} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # GB300 exposed time vs the FlashInfer prequantized GEMM (behind a PDL predecessor), at 24 / 48 tokens: QKV 3.76 / + # 4.63 vs 7.46 / 11.44 us, o_proj 2.74 / 3.33 vs 4.46 / 5.84, shared gate_up 4.57 / 5.71 vs 7.16 / 11.25. Behind + # only for K = 768 (unsplit shared down) at 17-32 tokens: 3.1 vs 2.9. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], + x_swiglu: bool = False, +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + if nacc != 1: + # More than one TMEM accumulator is broken on SM107 (783556 / 783557 / 783579: nondeterministic or garbage + # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + x_swiglu, + ) + return fn + + +def _views( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight, + weight_scale, + y, + x_bf16: bool, + x_swiglu: bool = False, +) -> tuple: + # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. + k = x.shape[1] // 2 if x_swiglu else x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, + x_swiglu: bool = False, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + if x_swiglu: + k //= 2 + n = weight.shape[0] + if weight.shape[1] != k: + return None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu) + if resolved is None: + return None + key, runtime_ar = resolved + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, + x_swiglu: bool = False, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros( + m, + 2 * k if x_swiglu else k, + dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, + device=dev, + ) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_swiglu( + gate_up: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``m3_mx_gemv_bf16(swiglu(gate_up), ...)`` in one launch: ``gate_up`` bf16 [M <= 8, 2K] (contiguous, gate then + up) is activated in the kernel's quantize prologue with the swigluoai arithmetic of the Triton silu_and_mul kernel + (alpha 1.702, limit 7, the (up + 1) offset; bit for bit), then quantized as the bf16 path does.""" + if ( + gate_up.dtype != torch.bfloat16 + or gate_up.dim() != 2 + or not gate_up.is_contiguous() + or gate_up.data_ptr() % 16 != 0 + or gate_up.shape[1] % 2 != 0 + ): + return None + return _run(gate_up, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar, True) + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/late_compile.py b/tensorrt_llm/_torch/late_compile.py new file mode 100644 index 000000000000..97eb4bf210ae --- /dev/null +++ b/tensorrt_llm/_torch/late_compile.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Late-compile detector: report kernel compiles that happen after the executor finished its warmup. + +A compile inside a served request costs that request its whole compile time, and in-memory compiles (CuTe DSL +handles, Triton specializations) leave no file in the JIT cache directories, so cache-file checks cannot see them. +After ``mark_server_ready()`` every ``cutlass.cute.compile`` call and every Triton JIT compile logs one warning line +``[late-compile] op= key=``. The executor arms it only for a model whose RuntimeHooks set +late_compile_detect. Kill switch: TRTLLM_LATE_COMPILE_DETECT=0. + +Not covered here: the TRTLLM-Gen FMHA runtime NVRTC compile (C++, an in-memory module cache warmed for generation +kernels only), whose warmup and persistence fix is tracked with that kernel library, and the fmha_sm100 package's +on-disk AOT cache, whose post-warmup "[aot_cache] Saved" lines are its own signal. +""" + +import os +from typing import Any + +from tensorrt_llm.logger import logger + +_ENABLED = os.environ.get("TRTLLM_LATE_COMPILE_DETECT", "1") != "0" +_ready = [False] + + +def note(op: str, key: Any) -> None: + """Log a compile of ``op`` for ``key`` when it happens after warmup.""" + if _ready[0]: + logger.warning(f"[late-compile] op={op} key={key}") + + +class _CuteCompile: + """Stands in for ``cutlass.cute.compile``: notes each call, then forwards it (``cute.compile[...]`` included).""" + + def __init__(self, inner: Any) -> None: + self._inner = inner + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + target = args[0] if args else None + note("cute.compile", getattr(target, "__qualname__", type(target).__name__)) + return self._inner(*args, **kwargs) + + def __getitem__(self, item: Any) -> "_CuteCompile": + return _CuteCompile(self._inner[item]) + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + +def _hook_cute() -> None: + try: + import cutlass.cute as cute + except ImportError: + return + if not isinstance(cute.compile, _CuteCompile): + cute.compile = _CuteCompile(cute.compile) + + +def _hook_triton() -> None: + try: + import triton + except ImportError: + return + runtime = getattr(getattr(triton, "knobs", None), "runtime", None) + if runtime is None or not hasattr(runtime, "jit_post_compile_hook"): + logger.info( + "[late-compile] Triton post-compile hook unavailable; Triton compiles are not reported" + ) + return + previous = runtime.jit_post_compile_hook + + def hook(*args: Any, **kwargs: Any) -> Any: + try: + fn = kwargs.get("fn") + name = getattr(getattr(fn, "fn", fn), "__qualname__", "?") + note("triton", f"{name} {kwargs.get('key', '')}") + except Exception: # noqa: BLE001 (the detector must never break a compile) + pass + return previous(*args, **kwargs) if previous is not None else None + + runtime.jit_post_compile_hook = hook + + +def mark_server_ready() -> None: + """Arm the detector once the executor's warmup is done (idempotent).""" + if not _ENABLED or _ready[0]: + return + _hook_cute() + _hook_triton() + _ready[0] = True + logger.info("[late-compile] detector armed after warmup") diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index e1c3c71bcd73..f08941905697 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -4611,9 +4611,16 @@ def instantiate_sampler( if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder( ): sampler = get_spec_decoder(sampler_args, engine.spec_config) - spec_step_fold = runtime_hooks_of(engine.model).spec_step_fold - if spec_step_fold is not None and isinstance(sampler, SpecSampler): - sampler.bind_spec_step_fold(spec_step_fold) + runtime_hooks = runtime_hooks_of(engine.model) + if isinstance(sampler, SpecSampler): + if runtime_hooks.spec_step_fold is not None: + sampler.bind_spec_step_fold(runtime_hooks.spec_step_fold) + sampler.d2h_side_stream_copies = runtime_hooks.sampler_d2h_side_stream + sampler.static_slot_ids = runtime_hooks.spec_small_h2d + if runtime_hooks.sampler_d2h_side_stream: + logger.info( + "[M3 host] spec sampler store copies on the D2H side stream" + ) return sampler if mm_encoder_only: diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 2e131a2a246b..85ff65d396d2 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -205,6 +205,8 @@ def __init__(self, config: CUDAGraphRunnerConfig): self._pre_replay_position_ids: Optional[torch.Tensor] = None self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} + # Graphs captured with a step-input pack's prologue, and that pack (step_inputs.capture_prologue). + self._step_input_prologue_keys: Dict[KeyType, Any] = {} self.graph_outputs: Dict[KeyType, Callable[[], Optional[torch.Tensor]]] = {} self.graph_metadata: Dict[KeyType, Dict[str, Any]] = {} @@ -792,6 +794,11 @@ def capture(self, # setup/capture; release its reference before entering. output = None with torch.cuda.graph(graph, pool=self.memory_pool): + # A pre-replay hook reads the static inputs before the replay, so their copies must land before it. + prologue_pack = (None if self._pre_replay_hooks else + step_inputs.capture_prologue()) + if prologue_pack is not None: + self._step_input_prologue_keys[key] = prologue_pack output = forward_fn(capture_inputs) # Snapshot before postprocessing so the check compares against # the addresses the graph actually baked in. @@ -910,8 +917,14 @@ def replay( static_encoder_hidden_states[actual_num_encoder_tokens:].zero_() # Issue the step's packed input copies (no-op without an open pack) before - # anything below reads the static inputs on the device. - step_inputs.flush() + # anything below reads the static inputs on the device; a graph captured + # with the pack's prologue issues them itself at the start of the replay. + prologue_pack = self._step_input_prologue_keys.get(key) + step_inputs.flush(in_graph=prologue_pack is not None) + if prologue_pack is not None and not step_inputs.is_active( + prologue_pack): + # Replayed outside the pack's context: nothing filled its window. + prologue_pack.quiesce_graph_window() if self._pre_replay_hooks: hook_tensors = static_tensors if position_id_offsets is not None: @@ -1363,6 +1376,7 @@ def clear(self): self.graph_outputs.clear() self.graph_metadata.clear() self._pre_replay_position_ids = None + self._step_input_prologue_keys.clear() self.padding_dummy_requests = {} del self.memory_pool self.memory_pool = None diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index d134a9d5155f..465510f00873 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -1730,9 +1730,11 @@ def previous_seq_slots_device() -> torch.Tensor: num_accepted_draft_tokens = torch.tensor( num_accepted_draft_tokens, dtype=torch.int, pin_memory=prefer_pinned() ) + # Read only by the forward (through the spec metadata), so the copy may wait for the pack's flush. copy_h2d_if_changed( self._buffers.num_accepted_draft_tokens_cuda[: len(num_accepted_draft_tokens)], num_accepted_draft_tokens, + pack=self.runtime_hooks.spec_small_h2d, ) step_fold = self.runtime_hooks.spec_step_fold if self._config.is_spec_decode else None fold_rows = None @@ -2313,6 +2315,20 @@ def _get_remote_tail_decode_rank_tokens( ) return gathered[:, 0].tolist() + def _kv_lens_corrected_in_preprocess( + self, inputs: dict[str, Any], enable_spec_decode: bool + ) -> bool: + """Whether preprocess_inputs applies the overlap scheduler's KV-length offsets to kv_lens_cuda and then calls + on_update_kv_lens (the branch below).""" + attn_meta = inputs.get("attn_metadata") + return ( + enable_spec_decode + and not self._config.disable_overlap_scheduler + and attn_meta is not None + and attn_meta.kv_cache_manager is not None + and hasattr(attn_meta, "kv_lens_cuda") + ) + def preprocess_inputs( self, inputs: dict[str, Any], *, enable_spec_decode: bool, runtime_draft_len: int ) -> dict[str, Any]: @@ -2320,8 +2336,14 @@ def preprocess_inputs( Make some changes to the device inputs and avoid blocking the async data transfer """ attn_meta = inputs.get("attn_metadata") - # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. - if attn_meta is not None: + # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. With + # RuntimeHooks.kv_lens_update_once, a step that corrects the KV lengths below (and then calls it again) calls + # it only there: the second call rewrites everything the first one writes, and nothing reads the metadata + # in between. + if attn_meta is not None and not ( + self.runtime_hooks.kv_lens_update_once + and self._kv_lens_corrected_in_preprocess(inputs, enable_spec_decode) + ): attn_meta.on_update_kv_lens() if enable_spec_decode and not self._config.disable_overlap_scheduler: @@ -2338,7 +2360,25 @@ def preprocess_inputs( num_ctx_tokens = inputs["attn_metadata"].num_ctx_tokens num_chunked_ctx_requests = inputs["attn_metadata"].num_chunked_ctx_requests previous_batch_tokens = inputs["input_ids"].shape[0] - num_ctx_tokens - if inputs["position_ids"].ndim == 3: # mrope: [3, 1, N] + # RuntimeHooks.spec_offsets_op: both corrections of a generation-row step in one launch. + fused_offsets = ( + self.runtime_hooks.spec_offsets_op is not None + and inputs["position_ids"].ndim != 3 + and hasattr(inputs["attn_metadata"], "kv_lens_cuda") + and not ( + num_ctx_requests >= num_chunked_ctx_requests + and num_chunked_ctx_requests > 0 + ) + and self.runtime_hooks.spec_offsets_op( + inputs["position_ids"][0, num_ctx_tokens:], + self._buffers.previous_pos_id_offsets_cuda[:previous_batch_tokens], + inputs["attn_metadata"].kv_lens_cuda[num_ctx_requests:num_seqs], + self._buffers.previous_kv_lens_offsets_cuda[:num_gen_requests], + ) + ) + if fused_offsets: + pass + elif inputs["position_ids"].ndim == 3: # mrope: [3, 1, N] inputs["position_ids"][:, :, num_ctx_tokens:] += ( self._buffers.previous_pos_id_offsets_cuda[:previous_batch_tokens] ) @@ -2348,7 +2388,9 @@ def preprocess_inputs( ) if hasattr(inputs["attn_metadata"], "kv_lens_cuda"): - if ( + if fused_offsets: + pass + elif ( num_ctx_requests >= num_chunked_ctx_requests and num_chunked_ctx_requests > 0 ): diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py index 866c16adbcfe..7ed0c04791e1 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py @@ -25,6 +25,7 @@ runtime_hooks_of, ) from tensorrt_llm._torch.speculative import SpecMetadata, get_spec_metadata +from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from ...metadata import build_attention_metadata @@ -135,6 +136,7 @@ def _arm_msa_policy(self, kv_cache_manager: KVCacheManager | KVCacheManagerV2 | kv_cache_manager.configure_decode_shortcuts( skip_unchanged_block_offsets=msa.skip_unchanged_block_offsets, reuse_pool_views=msa.reuse_pool_views, + block_ids_from_views=msa.block_ids_from_views, ) def setup_speculative( @@ -146,6 +148,13 @@ def setup_speculative( self._config, self.model.config, spec_resource_manager ) self.spec_metadata.skip_greedy_rng_upload = self.runtime_hooks.skip_greedy_rng_upload + self.spec_metadata.pack_rng_uploads = self.runtime_hooks.pack_spec_rng_uploads + if self.runtime_hooks.pack_spec_rng_uploads: + logger.info("[M3 host] sampled-step Philox uploads staged into the step-input pack") + if self.runtime_hooks.spec_small_h2d: + logger.info( + "[M3 host] spec step small uploads off the boundary (accepted counts packed, static slot ids)" + ) return self.spec_metadata diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 05df4e1b8396..f1dde6db9dea 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -63,6 +63,7 @@ PrefillCudaGraphBackend, SeqLenAwareSparseAttentionConfig, ) +from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from ...lora import LoraParamBuilder, make_cuda_graph_lora_manager @@ -200,12 +201,41 @@ def __init__( self._config = config self.runtime_hooks = runtime_hooks_of(model) - self._h2d_cache = H2DSkipCache() if self.runtime_hooks.skip_unchanged_h2d else None + self._h2d_cache = ( + H2DSkipCache(changed_span=self.runtime_hooks.h2d_changed_span) + if self.runtime_hooks.skip_unchanged_h2d + else None + ) self._step_pack = ( - StepInputPack(torch.device("cuda", torch.cuda.current_device())) + StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=self.runtime_hooks.pack_parallel_scatter, + in_graph=self.runtime_hooks.pack_in_graph, + early_copy=self.runtime_hooks.pack_early_copy, + step_epoch=self.runtime_hooks.pack_step_epoch, + prologue_op=self.runtime_hooks.pack_prologue_op, + ) if self.runtime_hooks.pack_step_inputs else None ) + if self._step_pack is not None and self._step_pack.parallel_scatter: + logger.info("[M3 host] step-input pack: parallel scatter for non-overlapping copies") + if self._step_pack is not None and self._step_pack.in_graph: + logger.info("[M3 host] step-input pack issued by the decode graph's prologue") + if self._step_pack is not None and self._step_pack.early_copy: + logger.info( + "[M3 host] step-input pack: window uploaded ahead of the replay on a side copy stream" + ) + if self._step_pack is not None and self._step_pack.epoch is not None: + logger.info( + "[M3 host] step epoch: one device counter bumped by the decode graph's prologue" + ) + if self._step_pack is not None and self._step_pack.prologue_op is not None: + logger.info( + "[M3 host] step-input pack: the decode graph's prologue is one launch (scatter, header, epoch)" + ) + if self._h2d_cache is not None and self._h2d_cache.changed_span: + logger.info("[M3 host] changed metadata uploads stage only each row's changed span") self.forward_pass_callable = None # Optional in-graph sampling hook, registered by PyExecutor. Called at @@ -491,7 +521,11 @@ def _short_long_seq_len_threshold(self) -> int | None: Set by a sparse attention config that asks for the split, or by a model whose MSA policy stages identity selections for short sequences: a sequence within the - top-k blocks attends every valid block. + top-k blocks attends every valid block. The short family runs the identity selection, + so every token it attends must fit in the top-k blocks. The graph key estimates a + sequence's length from host token counts, which under the overlap scheduler miss up to + the in-flight step's accepted tokens; the identity threshold leaves one generation + step's tokens of headroom for that. """ sparse_config = self.sparse_attention_config if ( @@ -505,7 +539,8 @@ def _short_long_seq_len_threshold(self) -> int | None: "MsaRuntimePolicy.short_sequence_identity requires a MiniMax-M3 sparse " f"attention config, got {type(sparse_config).__name__}" ) - return sparse_config.sparse_topk_blocks * sparse_config.sparse_block_size + topk_span = sparse_config.sparse_topk_blocks * sparse_config.sparse_block_size + return topk_span - self.get_runtime_tokens_per_gen_step(self._config.max_draft_len) return None def _initialize_breakable_cuda_graph_runner( diff --git a/tensorrt_llm/_torch/pyexecutor/h2d_cache.py b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py index 1695ac2e1189..b8827c3b3c81 100644 --- a/tensorrt_llm/_torch/pyexecutor/h2d_cache.py +++ b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py @@ -14,10 +14,11 @@ cache travels in the context rather than as an argument. Route a copy through it only for a destination that no other code writes (no device kernel, no other copy path, no other -engine): the shadow describes the device contents only under that invariant. Pooled metadata buffers satisfy it because -an engine that arms the skip gives its CUDA graph metadata a buffer pool of its own -(``RuntimeHooks.skips_unchanged_metadata_uploads``); eager metadata allocates its own buffers. A write through a -different view of the same buffer replaces the shadow, so overlapping views never skip on stale state. +engine, no graph replay): the shadow describes the device contents only under that invariant. CUDA graph metadata +buffers do not satisfy it: every graph key's metadata takes them from one pool, which hands the same block to each key +that asks for the same name, and a captured graph's device writes into a block run again at each of its replays. Those +destinations use ``copy_h2d_pool_shared`` / ``fill_tail_pool_shared``, which never skip. A write through a different +view of the same buffer replaces the shadow, so overlapping views never skip on stale state. """ from __future__ import annotations @@ -32,9 +33,19 @@ class H2DSkipCache: - """One engine's host shadows of its skip-routed destinations (see the module docstring).""" + """One engine's host shadows of its skip-routed destinations (see the module docstring). - def __init__(self) -> None: + ``changed_span``: a packed copy into a 2-D view with contiguous rows whose shadow covers the same rows (same + first element and row stride) uploads, per row, only the span from its first to its last changed column, plus any + columns the shadow does not cover. A page table that gained one page then stages a few bytes instead of every + padded row, so the pack still fits the decode graph's prologue window. + """ + + # A changed-span upload is used only while it moves at most this fraction of the view's bytes. + CHANGED_SPAN_MAX_FRACTION = 0.5 + + def __init__(self, changed_span: bool = False) -> None: + self.changed_span = changed_span # Device storage address -> (the storage, kept alive so the address is not reused; the view key; the host # values). Keyed by storage rather than by view base: tensors made under inference mode do not record their # base. @@ -65,13 +76,58 @@ def copy(self, dst: torch.Tensor, src: torch.Tensor, pack: bool) -> bool: and torch.equal(entry[2], src) ): return False - if pack: - step_inputs.copy_(dst, src) - else: - dst.copy_(src, non_blocking=True) + if not ( + pack + and self.changed_span + and entry is not None + and self._copy_changed_span(dst, src, entry[1:]) + ): + if pack: + step_inputs.copy_(dst, src) + else: + dst.copy_(src, non_blocking=True) self._shadow[addr] = (storage, key, src.clone()) return True + def _copy_changed_span(self, dst: torch.Tensor, src: torch.Tensor, shadow: tuple) -> bool: + """Stage only the changed span of each row of ``dst`` (see the class docstring); False if not applicable.""" + (offset, shape, stride, dtype), old = shadow + if ( + dst.dim() != 2 + or dst.stride(1) != 1 + or offset != dst.storage_offset() + or len(shape) != 2 + or shape[0] != dst.shape[0] + or stride != tuple(dst.stride()) + or dtype != dst.dtype + or old.dtype != src.dtype + or tuple(src.shape) != tuple(dst.shape) + or not step_inputs.is_open() + ): + return False + cols = dst.shape[1] + known = min(shape[1], cols) + changed = src[:, :known] != old[:, :known] + if cols > known: + # Columns the shadow does not cover hold unknown device values. + changed = torch.cat([changed, changed.new_ones(dst.shape[0], cols - known)], dim=1) + rows = changed.any(dim=1).nonzero().flatten().tolist() + spans = [] + for row in rows: + columns = changed[row].nonzero().flatten() + spans.append((row, int(columns[0]), int(columns[-1]) + 1)) + if sum(hi - lo for _, lo, hi in spans) > self.CHANGED_SPAN_MAX_FRACTION * dst.numel(): + return False + for row, lo, hi in spans: + step_inputs.copy_(dst[row, lo:hi], src[row, lo:hi]) + return True + + def forget(self, dst: torch.Tensor) -> None: + """Drop the shadow and tail record of ``dst``'s storage (it is about to be written outside the cache).""" + addr = dst.untyped_storage().data_ptr() + self._shadow.pop(addr, None) + self._tails.pop(addr, None) + def fill_tail(self, buf: torch.Tensor, start: int, value: int) -> None: storage = buf.untyped_storage() addr = storage.data_ptr() @@ -93,6 +149,10 @@ def fill_tail(self, buf: torch.Tensor, start: int, value: int) -> None: _active: ContextVar[Optional[H2DSkipCache]] = ContextVar("trtllm_h2d_skip_cache", default=None) +# Kill switch for A/B only: True routes the pool-shared metadata destinations (``copy_h2d_pool_shared``, +# ``fill_tail_pool_shared``) through the skip cache again, as before. False keeps them always written. +SKIP_POOL_SHARED_UPLOADS = False + @contextlib.contextmanager def active(cache: Optional[H2DSkipCache]) -> Iterator[None]: @@ -135,3 +195,40 @@ def fill_tail_if_changed(buf: torch.Tensor, start: int, value: int) -> None: buf[start:].fill_(value) return cache.fill_tail(buf, start, value) + + +def copy_h2d_pool_shared(dst: torch.Tensor, src: torch.Tensor, pack: bool = False) -> bool: + """``copy_h2d_if_changed`` for a destination in a CUDA graph metadata buffer pool: always issued. + + Every graph key's metadata takes its buffers from one pool, which hands the same block to each key that asks + for the same buffer name, and the captured graphs write some of those blocks on the device at every replay. The + host shadow therefore does not describe the device contents, so such a copy is never skipped. ``pack``: as for + ``copy_h2d_if_changed``, the copy may join the open ``step_inputs`` pack when a cache is active. + """ + cache = _active.get() + if cache is None: + dst.copy_(src, non_blocking=True) + return True + if SKIP_POOL_SHARED_UPLOADS: + return cache.copy(dst, src, pack) + cache.forget(dst) + if pack: + step_inputs.copy_(dst, src) + else: + dst.copy_(src, non_blocking=True) + return True + + +def fill_tail_pool_shared(buf: torch.Tensor, start: int, value: int) -> None: + """``buf[start:].fill_(value)`` for a buffer in a CUDA graph metadata buffer pool: always issued. + + See ``copy_h2d_pool_shared``: another graph key's replay can write ``buf`` past ``start``. + """ + cache = _active.get() + if cache is not None and SKIP_POOL_SHARED_UPLOADS: + cache.fill_tail(buf, start, value) + return + if cache is not None: + cache.forget(buf) + if start < buf.shape[0]: + buf[start:].fill_(value) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9f2f4e12a855..e11d3f87d335 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -31,6 +31,7 @@ from tensorrt_llm._utils import get_sm_version, str_dtype_to_torch from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, ExecutorMemoryType, + MiniMaxM3SparseAttentionConfig, ModelExpressConfig, SparseAttentionConfig, TorchLlmArgs) from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, @@ -165,9 +166,16 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, effective_kv_cache_quant = (kv_cache_quant if pyt_kv_cache_dtype == "auto" else mapped_pyt_quant) + # MiniMax-M3 MSA keeps NVFP4 only on its sparse layers, which MSA and + # Triton serve; its dense and draft layers stay FP8 for trtllm-gen. + sparse_config = model_config.sparse_attention_config + uses_minimax_m3_msa = (isinstance(sparse_config, + MiniMaxM3SparseAttentionConfig) + and sparse_config.implementation == "msa") if (effective_kv_cache_quant in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value) and not supports_fp4_mla_attention(model_config) - and torch.cuda.is_available() and get_sm_version() == 107): + and not uses_minimax_m3_msa and torch.cuda.is_available() + and get_sm_version() == 107): logger.warning( "NVFP4 KV cache is not supported by trtllm-gen on SM107; " "using FP8 KV cache instead.") diff --git a/tensorrt_llm/_torch/pyexecutor/profiling.py b/tensorrt_llm/_torch/pyexecutor/profiling.py index e7bcbd5157fd..15ccb51c6094 100644 --- a/tensorrt_llm/_torch/pyexecutor/profiling.py +++ b/tensorrt_llm/_torch/pyexecutor/profiling.py @@ -55,6 +55,8 @@ from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from tensorrt_llm.tools.profiler.host_profile_tools.host_profiler import get_global_profiler +from .runtime_hooks import runtime_hooks_of + # Environment-variable names used by both the env-var-driven legacy # path and the runtime HTTP-endpoint path. They are duplicated as # module-level constants here so callers that only need the manager @@ -608,6 +610,18 @@ def profile_step(self): start_event_2 = None end_event_2 = torch.cuda.Event(enable_timing=True) prev_device_step_time = None + # With RuntimeHooks.iteration_timing_one_event, boundary_events[k % 3] marks the boundary before + # iteration k, recorded on timing_stream, and replaces the start/end ping-pong above. + one_timing_event = runtime_hooks_of( + getattr(executor.model_engine, "model", None) + ).iteration_timing_one_event + boundary_events = ( + [torch.cuda.Event(enable_timing=True) for _ in range(3)] if one_timing_event else None + ) + boundary_count = 0 + timing_stream = torch.cuda.Stream() if one_timing_event else None + if one_timing_event: + logger.info("[M3 host] iteration timing: one side-stream event per iteration boundary") # Cumulative fetch counter as of the previous iter log line; the # per-iteration ``num_fetched`` on the log line is the delta. prev_fetch_requests_cur_rank = 0 @@ -663,6 +677,7 @@ def profile_step_fn(): nonlocal it, enabled, start_time nonlocal start_event_1, end_event_1, start_event_2, end_event_2 nonlocal prev_device_step_time, prev_fetch_requests_cur_rank + nonlocal boundary_count nonlocal torch_profiler, active_torch_trace_path nonlocal active_enable_torch_trace calibrator.post_step(it) @@ -718,9 +733,21 @@ def profile_step_fn(): # ``self`` lets the /metrics serializer pick up the # values without going through the log line. should_capture_timing = executor.print_log or executor.enable_iter_perf_stats + if should_capture_timing and one_timing_event: + # Iteration k - 2 spans the two previous boundaries: the step the ping-pong reports here. + timing_stream.wait_stream(torch.cuda.current_stream()) + boundary_events[boundary_count % 3].record(timing_stream) + if boundary_count >= 2: + boundary_events[(boundary_count - 1) % 3].synchronize() + prev_device_step_time = boundary_events[(boundary_count - 2) % 3].elapsed_time( + boundary_events[(boundary_count - 1) % 3] + ) + boundary_count += 1 if should_capture_timing and start_time is not None: end_time = time.time() - if it % 2 == 0: + if one_timing_event: + pass + elif it % 2 == 0: end_event_1.record() if start_event_2 is not None: end_event_2.synchronize() @@ -847,7 +874,7 @@ def profile_step_fn(): calibrator.pre_step(it) start_time = time.time() - if should_capture_timing: + if should_capture_timing and not one_timing_event: if it % 2 == 0: if start_event_1 is None: start_event_1 = torch.cuda.Event(enable_timing=True) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9249685307c4..f0c83fa9cdb1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -114,7 +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 .runtime_hooks import RuntimeHooks, runtime_hooks_of from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors) from .scheduler import (RequestScheduler, ScheduledRequests, @@ -405,6 +405,21 @@ class PendingEncoderStep: result: Optional[EncoderStepResult] = None +def _arm_runtime_opt_ins(hooks: RuntimeHooks) -> bool: + """Arm the process-wide features ``hooks`` opt into; return whether to arm + the late-compile detector after warmup. + + Both are off for a model that declares neither: the SM107 MXFP8 GEMM + tactic table (the GEMM runner does not know the model, so it is armed + here) and the late-compile detector. + """ + if hooks.mxfp8_pinned_tactics: + # Imported on use: only an opted-in model needs the table. + from tensorrt_llm._torch.custom_ops import mxfp8_cute_pinned_tactics + mxfp8_cute_pinned_tactics.arm() + return hooks.late_compile_detect + + class PyExecutor: # True while the overlap loop runs with the execution stream current # (_overlap_loop_stream_scope). @@ -597,6 +612,18 @@ def __init__( self._pending_num_new_active_requests = 0 self._pending_num_completed_requests = 0 self.stream_interval = self.llm_args.stream_interval + # Streaming: also emit on decoding iterations 2..N when the model's + # runtime hooks ask for it (RuntimeHooks.first_content_emit_iters; 1, + # the default, keeps the stock schedule). + model_hooks = runtime_hooks_of(model_engine.model) + self.first_content_emit_iters = max( + 1, model_hooks.first_content_emit_iters) + if self.stream_interval > 1 and self.first_content_emit_iters > 1: + logger.info( + f"first-content emission engaged: emit on decoding iters " + f"<= {self.first_content_emit_iters} (stream_interval " + f"{self.stream_interval}), rank {self.dist.rank}") + self._late_compile_detect = _arm_runtime_opt_ins(model_hooks) self.perf_manager = PerfMetricsManager( enabled=getattr(self.llm_args, 'return_perf_metrics', False)) self.attention_dp_enable_balance = ( @@ -939,6 +966,9 @@ def __init__( self._warmup_encoder_cuda_graphs_enc_dec).result() self.is_warmup = False + if self._late_compile_detect: + from tensorrt_llm._torch.late_compile import mark_server_ready + mark_server_ready() # Snapshot some cumulative KV cache counters so that stats reported to # users exclude blocks reused and missed during warmup dummy requests. @@ -9219,6 +9249,24 @@ def _emit_first_token_responses(self, prev_scheduled_requests): self._enqueue_responses(new_responses) @nvtx_range("_handle_responses") + def _should_emit_response(self, + request: LlmRequest, + emit_first_iter: bool = True) -> bool: + """Whether ``_handle_responses`` emits a response for ``request`` now. + + Emits on decoding iteration 1, on finish, every ``stream_interval`` + iterations, and on iterations ``2..first_content_emit_iters``. + """ + decoding_iter = request.py_decoding_iter + if request.is_finished: + return True + if decoding_iter == 1: + # The early-emit prototype issues the (non-terminal) iter-1 + # response from `_emit_first_token_responses`; suppress it here. + return emit_first_iter + return (decoding_iter % self.stream_interval == 0 + or 1 < decoding_iter <= self.first_content_emit_iters) + def _handle_responses(self, emit_first_iter: bool = True): new_responses = [] requests_to_terminate = [] @@ -9310,15 +9358,7 @@ def _handle_responses(self, emit_first_iter: bool = True): "route_capture", None) if route_capture is not None: route_capture.attach_routes(request) # R3: append routes - should_emit = (request.py_decoding_iter == 1 or request.is_finished - or request.py_decoding_iter % self.stream_interval - == 0) - # The early-emit prototype issues the (non-terminal) iter-1 - # response from `_emit_first_token_responses`; suppress it here. - if (not emit_first_iter and request.py_decoding_iter == 1 - and not request.is_finished): - should_emit = False - if should_emit: + if self._should_emit_response(request, emit_first_iter): if request.return_perf_metrics: # Response creation may finalize and copy scalar ctx GPU totals. self.perf_manager.compute_batch_gpu_times( diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py index 8fa36c484922..3f0201f5f4de 100644 --- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py +++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py @@ -15,7 +15,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Optional, Protocol +from typing import TYPE_CHECKING, Callable, Optional, Protocol if TYPE_CHECKING: import torch @@ -144,6 +144,9 @@ class MsaRuntimePolicy: skip_unchanged_block_offsets: bool = False # Build each KV pool view once per pool geometry instead of on every attention call. reuse_pool_views: bool = False + # Fill the per-request block-id table by slice-copying each request's page-index view instead of + # building a Python list of the whole padded row (max_seq_len / tokens_per_block entries) first. + block_ids_from_views: bool = False @dataclass(frozen=True) @@ -153,19 +156,76 @@ class RuntimeHooks: # Pack a decode step's small host-to-device input copies into one upload and # one scatter launch, issued just before the CUDA graph replay. pack_step_inputs: bool = False + # Issue the step-input pack's host copies with one parallel scatter (a program per copy) instead of the + # serial one-CTA scatter, when their destinations do not overlap (pack_step_inputs). + pack_parallel_scatter: bool = False + # Issue the step-input pack inside the replayed graph (a prologue copies a pinned window and scatters it), so no + # operation is left between the decode graphs (needs pack_parallel_scatter). + pack_in_graph: bool = False + # With pack_in_graph: a side copy stream uploads the pinned window into the prologue's own device buffer while the + # previous graph still runs, so the prologue only waits on that copy and scatters device to device (no host copy + # at the head of the replay). + pack_early_copy: bool = False + # With pack_in_graph: the pack keeps a step epoch, one int64 device counter the graph's prologue increments once per + # replay and never resets, so in-graph per-step flags can be tagged with it instead of cleared between steps. + pack_step_epoch: bool = False + # With pack_early_copy: the model's one-launch prologue, called inside the capture as + # ``op(window, epoch_or_None, max_desc)``. It writes the window's host copies, then its device copies, zeroes the + # window's header and increments the epoch, which the default prologue does with two scatter launches and a fill. + pack_prologue_op: Optional[Callable[[torch.Tensor, Optional[torch.Tensor], int], None]] = None # Skip host-to-device metadata copies whose values the destination already holds. skip_unchanged_h2d: bool = False + # With skip_unchanged_h2d: a changed packed 2-D upload stages only each row's changed column span (a page table that + # gained a page re-sends a few bytes, not every padded row), so the pack keeps fitting the in-graph window. + h2d_changed_span: bool = False + # A speculative step under the overlap scheduler re-derives its KV-length metadata (on_update_kv_lens) once, + # after the forward's KV-length correction, instead of also before it: the first pass's device writes are all + # rewritten by the second. + kv_lens_update_once: bool = False + # The overlap scheduler's speculative step correction (position_ids += position offsets, kv_lens += KV-length + # offsets) as the model's one launch, called as ``op(positions, position_offsets, kv_lens, kv_lens_offsets)`` on the + # generation rows; it returns False to decline (the two in-tree adds run). + spec_offsets_op: Optional[ + Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], bool] + ] = None # Run the overlap executor loop with the execution stream current, so the # loop's own stream work and the forward are ordered on one stream. loop_on_execution_stream: bool = False # After autotuning, keep the fastest profiled tactic per shape on every # tensor-parallel rank, so the ranks capture the same kernels. merge_tuned_tactics_across_ranks: bool = False + # Streaming: also emit a response on decoding iterations 2..N, so the first + # counted text reaches the client one step after the first token instead of + # one stream interval later (a reasoning parser holds a leading reasoning + # tag back, so the iteration-1 chunk can carry no content). Emission only; + # tokens are unchanged. 1 keeps the stock schedule. + first_content_emit_iters: int = 1 + # On SM107, run FlashInfer's CuTe DSL MXFP8 GEMM above 32 rows at the pinned tuned tactics of + # custom_ops/mxfp8_cute_pinned_tactics (bitwise the default tactic's output, faster). The GEMM runner does not + # know the model, so the executor arms the table process-wide when the model opts in. + mxfp8_pinned_tactics: bool = False + # After warmup, log every CuTe DSL / Triton compile as a late compile (tensorrt_llm/_torch/late_compile.py). + late_compile_detect: bool = False # Skip the per-step Philox seed and offset uploads on all-greedy # one-model speculative batches (the argmax path never reads them). skip_greedy_rng_upload: bool = False + # Stage the one-model speculative sampler's per-step Philox seed and offset + # uploads into the step-input pack (pack_step_inputs), so they ride its one + # upload instead of each being a copy the next graph waits for. + pack_spec_rng_uploads: bool = False + # The one-model speculative step's small per-step uploads: the accepted-draft + # counts join the step-input pack, and the spec sampler keeps its slot ids in + # a device buffer it re-uploads only when they change. + spec_small_h2d: bool = False # Speculative step-input staging folded into the forward (see SpecStepFold). spec_step_fold: Optional[SpecStepFold] = None + # Copy the one-model speculative sampler's per-slot store to the host on the sampler's D2H side + # stream, so those copies overlap the next step's input staging instead of preceding it. + sampler_d2h_side_stream: bool = False + # Time the iteration log's device step with one timing event per iteration boundary, recorded on a side + # stream that waits for the current stream, instead of an end and a start event back to back on the + # current stream (the same instant, with no GPU work between them). The reported step times are unchanged. + iteration_timing_one_event: bool = False # MiniMax-M3 MSA metadata preparation. msa: MsaRuntimePolicy = MsaRuntimePolicy() diff --git a/tensorrt_llm/_torch/pyexecutor/step_inputs.py b/tensorrt_llm/_torch/pyexecutor/step_inputs.py index ff0c02a95611..29b3eb318afd 100644 --- a/tensorrt_llm/_torch/pyexecutor/step_inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -19,15 +19,47 @@ destination between the copy site and the flush: the pack moves the copy to the flush point. A copy that cannot be packed (non-contiguous, device-side dtype change) flushes the pack first and then runs as a plain copy, which keeps the original order. A pack holds every tensor it staged until it is issued, so a temporary source outlives its copy. + +With ``parallel_scatter`` (``RuntimeHooks.pack_parallel_scatter``) a pack of host copies only, whose destinations do +not overlap, is issued with one program per copy (``_parallel_scatter_kernel``) instead of the serial scatter: the same +bytes land in the same places, and with no overlap their order cannot matter. Any other pack takes the serial scatter. + +With ``in_graph`` (``RuntimeHooks.pack_in_graph``, on top of ``parallel_scatter``) the issue itself moves into the +replayed graph: a graph captured with ``capture_prologue`` starts by copying a fixed pinned window (``ghost``) to the +device staging buffer, records ``ghost_free``, and runs the parallel scatter with the copy counts read from the staged +header (one program per descriptor slot; the slots past the counts write nothing). ``flush(in_graph=True)`` before +such a replay then only fills the window on the host, after ``ghost_free`` shows the previous replay has copied it: no +operation is left between the graphs. A pack that does not fit the window or is not order-free is issued eagerly as +before, with a zero-count window so the graph's prologue writes nothing. + +With ``early_copy`` (``RuntimeHooks.pack_early_copy``, on top of ``in_graph``) the window reaches the device before the +replay: ``flush(in_graph=True)`` fills the pinned window and enqueues its upload on a side copy stream, into a device +buffer of the prologue's own (``gdev``), while the previous graph may still run. The prologue then waits for that upload +(``staged_ready``) and scatters device to device; it zeroes the window's header after the scatter and records +``prologue_done``, which the next upload waits for before it overwrites ``gdev``. The next fill of the pinned window +waits for ``ghost_free``, which the upload records once it has read the window. + +With ``step_epoch`` (``RuntimeHooks.pack_step_epoch``, on top of ``in_graph``) the pack keeps one int64 device counter, +``StepInputPack.epoch`` (``step_epoch()`` for the active pack), which the prologue's first scatter launch increments +once per replay (program 0, before any other kernel of the graph runs). It is never reset: it starts at 0, advances once +when the prologue kernel is compiled and then once per replay of any graph with the prologue, with or without a new pack +to scatter, and not on an eager forward. A kernel later in the graph reads it after its programmatic-dependent-launch +wait, so per-step device flags can be tagged with it instead of being cleared between steps. + +With ``prologue_op`` (``RuntimeHooks.pack_prologue_op``, on top of ``early_copy``) the prologue is the model's one +launch, ``prologue_op(gdev, epoch, MAX_DESC)``, in place of the two scatter launches and the header fill: it writes the +same bytes, zeroes the same header and increments the epoch the same once per replay (and once at warm-up). """ from __future__ import annotations import contextlib from contextvars import ContextVar -from typing import Iterator, Optional +from typing import Callable, Iterator, Optional import torch +import triton +import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import ( # noqa: F401 fake_ops as _m3_fake_ops, # registers trtllm::step_inputs_scatter's fake impl @@ -35,6 +67,87 @@ from tensorrt_llm._utils import prefer_pinned +# One compiled variant for every descriptor offset: specializing ``first`` (== 1, divisible by 16) would compile a +# new kernel when a step's pack first holds a new number of copies, after the warm-up. +@triton.jit(do_not_specialize=["first"]) +def _parallel_scatter_kernel( + stage_u8, stage_i64, first, DEVICE_SRC: tl.constexpr, BLOCK: tl.constexpr +): + """Program i writes copy ``first + i`` ([dst, src, bytes] descriptors after the 2-word header). ``src`` is an offset + into the staged buffer for a host copy, a device address for a device copy (``DEVICE_SRC``).""" + d = 2 + 3 * (first + tl.program_id(0)) + dst = tl.load(stage_i64 + d).to(tl.pointer_type(tl.uint8)) + src_word = tl.load(stage_i64 + d + 1) + if DEVICE_SRC: + src = src_word.to(tl.pointer_type(tl.uint8)) + else: + src = stage_u8 + src_word + nbytes = tl.load(stage_i64 + d + 2) + for start in range(0, nbytes, BLOCK): + o = start + tl.arange(0, BLOCK) + tl.store(dst + o, tl.load(src + o, mask=o < nbytes), mask=o < nbytes) + + +@triton.jit +def _graph_scatter_kernel( + stage_u8, + stage_i64, + epoch, + DEVICE_SRC: tl.constexpr, + BLOCK: tl.constexpr, + BUMP_EPOCH: tl.constexpr, +): + """Descriptor slot i of the host copies (``DEVICE_SRC`` False) or of the device copies, with the counts read from + the staged header; a slot past its count writes nothing. ``BUMP_EPOCH``: program 0 also increments ``epoch[0]``.""" + i = tl.program_id(0) + if BUMP_EPOCH: + if i == 0: + tl.store(epoch, tl.load(epoch) + 1) + num_h2d = tl.load(stage_i64) + if DEVICE_SRC: + count = tl.load(stage_i64 + 1) + d = 2 + 3 * (num_h2d + i) + else: + count = num_h2d + d = 2 + 3 * i + if i < count: + dst = tl.load(stage_i64 + d).to(tl.pointer_type(tl.uint8)) + src_word = tl.load(stage_i64 + d + 1) + if DEVICE_SRC: + src = src_word.to(tl.pointer_type(tl.uint8)) + else: + src = stage_u8 + src_word + nbytes = tl.load(stage_i64 + d + 2) + for start in range(0, nbytes, BLOCK): + o = start + tl.arange(0, BLOCK) + tl.store(dst + o, tl.load(src + o, mask=o < nbytes), mask=o < nbytes) + + +def _disjoint(ranges: list[tuple[int, int]]) -> bool: + """Whether the [start, start + bytes) ranges are pairwise disjoint.""" + end = -1 + for start, nbytes in sorted(ranges): + if start < end: + return False + end = start + nbytes + return True + + +def _parallel_ok(h2d: list[tuple[int, int, int]], d2d: list[tuple[int, int, int]]) -> bool: + """Whether issuing the host copies in parallel and then the device copies in parallel writes what the serial + scatter (every copy in order) writes: the host copies' destinations are disjoint, and the device copies' + destinations are disjoint and none of them is a device copy's source (a host copy may still feed a device copy, + as the launches run in order).""" + if not _disjoint([(dst, n) for dst, _, n in h2d]): + return False + if not d2d: + return True + dsts = [(dst, n) for dst, _, n in d2d] + if not _disjoint(dsts): + return False + return all(_disjoint(dsts + [(src, n)]) for _, src, n in d2d) + + class StepInputPack: """One engine's pack of deferred step-input copies (see the module docstring).""" @@ -45,8 +158,47 @@ class StepInputPack: HEADER_BYTES = 16 DESC_BYTES = 24 ALIGN = 16 + # The pinned window a graph's prologue copies (in_graph): the header, the descriptors and up to ~10 KB of data. + GRAPH_WINDOW = 1 << 14 - def __init__(self, device: torch.device) -> None: + def __init__( + self, + device: torch.device, + parallel_scatter: bool = False, + in_graph: bool = False, + early_copy: bool = False, + step_epoch: bool = False, + prologue_op: Optional[Callable[[torch.Tensor, Optional[torch.Tensor], int], None]] = None, + ) -> None: + self.parallel_scatter = parallel_scatter + self.in_graph = in_graph and parallel_scatter + self.early_copy = early_copy and self.in_graph + # The model's one-launch prologue (see the module docstring); None: the scatter launches and the header fill. + self.prologue_op = prologue_op if self.early_copy else None + # The step epoch (see the module docstring); None without step_epoch. + self.epoch: Optional[torch.Tensor] = ( + torch.zeros(1, dtype=torch.int64, device=device) + if step_epoch and self.in_graph + else None + ) + if self.early_copy: + # The prologue's own device window: an eager flush (self.dev) can never leave a header in it. + self.gdev = torch.zeros(self.GRAPH_WINDOW, dtype=torch.uint8, device=device) + self.copy_stream = torch.cuda.Stream(device=device) + # Recorded by the side upload once it landed (the prologue waits on it) ... + self.staged_ready = torch.cuda.Event(external=True) + # ... and by every replayed prologue after its scatter (the next upload waits on it). + self.prologue_done = torch.cuda.Event(external=True) + # A CUDA event is created on its first record, and a wait on one never created is skipped: a capture + # would then leave the prologue's wait out of the graph. Create both now. + self.staged_ready.record(self.copy_stream) + self.prologue_done.record(self.copy_stream) + if self.in_graph: + self.ghost = torch.zeros( + self.GRAPH_WINDOW, dtype=torch.uint8, pin_memory=prefer_pinned() + ) + # Recorded by every replayed prologue right after it copied the window (an event node). + self.ghost_free = torch.cuda.Event(external=True) self.host = [ torch.empty(self.CAPACITY, dtype=torch.uint8, pin_memory=prefer_pinned()) for _ in range(self.RING) @@ -61,6 +213,86 @@ def __init__(self, device: torch.device) -> None: # Every tensor a staged copy reads or writes, kept alive until the pack is issued. self.staged: list[torch.Tensor] = [] self.used = self.data_start + if parallel_scatter: + self._warm_up(device) + if self.in_graph: + self._warm_up_graph_kernel() + + def _warm_up(self, device: torch.device) -> None: + """Compile the parallel scatter now, not on a serving step's flush.""" + dst = torch.zeros(8, dtype=torch.int32, device=device) + self.begin() + self.stage(dst[:4], torch.arange(4, dtype=torch.int32)) + self.stage(dst[4:], dst[:4]) + self.flush() + torch.cuda.synchronize(device) + + def _warm_up_graph_kernel(self) -> None: + """Compile the prologue's scatter now (a zero-count header: it writes nothing).""" + self.dev[: self.HEADER_BYTES].zero_() + if self.prologue_op is not None: + # The window it is captured with; its header is zero until the first upload. + self.prologue_op(self.gdev, self.epoch, self.MAX_DESC) + else: + self._launch_graph_scatter(self.dev) + torch.cuda.synchronize(self.dev.device) + + def _launch_graph_scatter(self, window: torch.Tensor) -> None: + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + bump = self.epoch is not None and not device_src + _graph_scatter_kernel[(self.MAX_DESC,)]( + window, + stage_i64, + self.epoch if bump else stage_i64, + DEVICE_SRC=device_src, + BLOCK=256, + BUMP_EPOCH=bump, + num_warps=1, + ) + + def capture_prologue(self) -> bool: + """Inside a graph capture: the prologue that issues the window ``flush(in_graph=True)`` fills (see the module + docstring). Returns whether the graph has it.""" + if not self.in_graph: + return False + if self.early_copy: + torch.cuda.current_stream().wait_event(self.staged_ready) + if self.prologue_op is not None: + self.prologue_op(self.gdev, self.epoch, self.MAX_DESC) + else: + self._launch_graph_scatter(self.gdev) + # Consumed: a replay with no new upload scatters nothing. + self.gdev[: self.HEADER_BYTES].zero_() + self.prologue_done.record() + return True + self.dev[: self.GRAPH_WINDOW].copy_(self.ghost, non_blocking=True) + self.ghost_free.record() + self._launch_graph_scatter(self.dev) + return True + + def _upload_window(self, nbytes: int) -> None: + """early_copy: upload the filled pinned window into gdev on the side stream, once the last replayed prologue + has consumed gdev; the window is fully written on the host before this is enqueued.""" + stream = self.copy_stream + stream.wait_event(self.prologue_done) + with torch.cuda.stream(stream): + self.gdev[:nbytes].copy_(self.ghost[:nbytes], non_blocking=True) + self.ghost_free.record(stream) + self.staged_ready.record(stream) + + def _ghost_writable(self) -> None: + # The last replayed prologue has copied the window (a no-op before the first replay). + self.ghost_free.synchronize() + + def quiesce_graph_window(self) -> None: + """Before replaying a graph with the prologue without a flush: a zero-count window, so it writes nothing.""" + if self.in_graph and self.ghost[: self.HEADER_BYTES].any(): + self._ghost_writable() + self.ghost[: self.HEADER_BYTES].zero_() + if self.early_copy: + # A window uploaded but never replayed would still sit in gdev: overwrite its header. + self._upload_window(self.HEADER_BYTES) def begin(self) -> None: if self.open: @@ -116,10 +348,20 @@ def stage(self, dst: torch.Tensor, src: torch.Tensor) -> bool: return True return False - def flush(self) -> None: + def flush(self, in_graph: bool = False) -> None: + """Issue the pack. ``in_graph``: the next replay is of a graph captured with ``capture_prologue``.""" if not self.open: + if in_graph: + self.quiesce_graph_window() return self.open = False + if in_graph and self.in_graph: + if self._fill_graph_window(): + if self.early_copy: + self._upload_window(self.used) + self._clear() + return + self.quiesce_graph_window() if not self.h2d and not self.d2d: return host = self.host[self.parity] @@ -129,7 +371,18 @@ def flush(self) -> None: host[: 8 * len(desc)].view(torch.int64).copy_(torch.tensor(desc, dtype=torch.int64)) # Descriptors and data in one upload; the gap between them is unused padding. self.dev[: self.used].copy_(host[: self.used], non_blocking=True) - torch.ops.trtllm.step_inputs_scatter(self.dev) + if self.parallel_scatter and _parallel_ok(self.h2d, self.d2d): + stage_i64 = self.dev.view(torch.int64) + if self.h2d: + _parallel_scatter_kernel[(len(self.h2d),)]( + self.dev, stage_i64, 0, DEVICE_SRC=False, BLOCK=256, num_warps=2 + ) + if self.d2d: + _parallel_scatter_kernel[(len(self.d2d),)]( + self.dev, stage_i64, len(self.h2d), DEVICE_SRC=True, BLOCK=256, num_warps=2 + ) + else: + torch.ops.trtllm.step_inputs_scatter(self.dev) event = self.events[self.parity] if event is None: event = self.events[self.parity] = torch.cuda.Event() @@ -138,6 +391,20 @@ def flush(self) -> None: # The scatter is queued on the current stream: the caching allocator keeps the staged tensors' memory for it. self._clear() + def _fill_graph_window(self) -> bool: + """Write the pack into the pinned window the replay's prologue issues; False when it does not fit or is not + order-free (the caller issues it eagerly).""" + if self.used > self.GRAPH_WINDOW or not _parallel_ok(self.h2d, self.d2d): + return False + desc = [len(self.h2d), len(self.d2d)] + [ + value for copy in self.h2d + self.d2d for value in copy + ] + self._ghost_writable() + host = self.host[self.parity] + self.ghost[: self.used].copy_(host[: self.used]) + self.ghost[: 8 * len(desc)].view(torch.int64).copy_(torch.tensor(desc, dtype=torch.int64)) + return True + def discard(self) -> None: """Close the pack without issuing it (its step failed).""" self.open = False @@ -164,6 +431,11 @@ def begin() -> None: pack.begin() +def is_active(pack: StepInputPack) -> bool: + """Whether ``pack`` is the calling context's pack.""" + return _active.get() is pack + + def is_open() -> bool: pack = _active.get() return pack is not None and pack.open @@ -190,8 +462,24 @@ def discard() -> None: pack.discard() -def flush() -> None: - """Issue the active open pack (one upload and one scatter launch on the current stream) and close it.""" +def flush(in_graph: bool = False) -> None: + """Issue the active open pack (one upload and one scatter launch on the current stream) and close it. + + ``in_graph``: the next replay is of a graph whose prologue issues the pack (``capture_prologue``).""" pack = _active.get() if pack is not None: - pack.flush() + pack.flush(in_graph=in_graph) + + +def step_epoch() -> Optional[torch.Tensor]: + """The active pack's step epoch (an int64 [1] device tensor), or None when there is none (see the module + docstring).""" + pack = _active.get() + return pack.epoch if pack is not None else None + + +def capture_prologue() -> Optional[StepInputPack]: + """Inside a graph capture: add the active pack's in-graph prologue; the pack whose window the graph issues, or + None when the graph has no prologue.""" + pack = _active.get() + return pack if pack is not None and pack.capture_prologue() else None diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index a32c1099340c..148ee4c550fa 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -32,6 +32,7 @@ from ..attention.backends.trtllm import (AttentionBackend, TrtllmAttention, TrtllmAttentionMetadata) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE +from ..pyexecutor import step_inputs from ..pyexecutor.resource_manager import ResourceManagerType if TYPE_CHECKING: @@ -575,6 +576,9 @@ def is_ragged_verify(self) -> bool: # from the model's RuntimeHooks). The argmax path never reads them, and the # next sampled batch rewrites every row it reads. skip_greedy_rng_upload: bool = False + # Stage the per-step Philox seed/offset uploads into the engine's step-input + # pack (RuntimeHooks.pack_spec_rng_uploads) instead of copying each one. + pack_rng_uploads: bool = False # Advanced-sampling specialization (deploy-time; from DecodingBaseConfig.advanced_sampling_mode). advanced_sampling_mode: AdvancedSamplingMode = AdvancedSamplingMode.FULL # Whether the occurrence penalties are enabled (deploy-time; from @@ -885,10 +889,15 @@ def _populate_request_rng_state( device='cuda') def _upload(dst: torch.Tensor, values: list[int]) -> None: - dst[:len(values)].copy_(torch.tensor(values, - dtype=torch.int64, - pin_memory=prefer_pinned()), - non_blocking=True) + src = torch.tensor(values, + dtype=torch.int64, + pin_memory=prefer_pinned()) + if self.pack_rng_uploads: + # The forward is the first reader of these buffers, so the copy + # may move to the pack's flush (a plain copy when no pack is open). + step_inputs.copy_(dst[:len(values)], src) + else: + dst[:len(values)].copy_(src, non_blocking=True) if self.skip_greedy_rng_upload and self.is_all_greedy_sample: # The offset window counters above still advanced. diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 8a51b3e69a35..db15c5fd189b 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -27,6 +27,7 @@ import torch +from ..._utils import prefer_pinned from ...sampling_params import SamplingParams from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..pyexecutor.resource_manager import BaseResourceManager @@ -73,6 +74,11 @@ "Drop top_p_decay / top_p_min from the request, or disable " "speculative decoding." ) +UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG = ( + "logits_processor (a logits post-processor callback) is not supported with one-model " + "speculative decoding. Drop logits_processor from the request, or disable speculative " + "decoding." +) UNSUPPORTED_MIN_P_MSG = ( "min_p requires 'advanced_sampling_mode: full' in the speculative decoding " "config when using one-model speculative decoding. Restore that default, " @@ -113,6 +119,10 @@ def one_model_sampling_rejection_reason( return UNSUPPORTED_NO_REPEAT_NGRAM_MSG if sampling_params.embedding_bias is not None: return UNSUPPORTED_EMBEDDING_BIAS_MSG + # A per-request callback becomes py_logits_post_processors; the one-model path never + # calls it (a fused verify may not even materialize the full logits). + if sampling_params.logits_processor: + return UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG # The same predicate top_p_decay_active() delegates to, so "active" cannot # mean one thing here and another in the sampler. if SamplingParams.params_imply_top_p_decay_active(sampling_params.top_p_decay): @@ -168,6 +178,13 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): # The model's speculative step-input staging, when it opted into one (bind_spec_step_fold). spec_step_fold: Optional[SpecStepFold] = None + # Copy the per-slot store to the host on the D2H side stream (RuntimeHooks.sampler_d2h_side_stream). + d2h_side_stream_copies: bool = False + # Keep the stored rows' slot ids in a device buffer re-uploaded only when they change + # (RuntimeHooks.spec_small_h2d). + static_slot_ids: bool = False + _slot_ids_host: Optional[list[int]] = None + _slot_ids_cuda: Optional[torch.Tensor] = None SampleState = SampleStateSpec @@ -232,8 +249,8 @@ def _validate_unsupported_logits_processors(request: LlmRequest) -> None: """Reject logits-side sampling features the one-model path cannot apply. TorchSampler implements min_length, bad_words, no_repeat_ngram_size, - embedding_bias and top_p_decay by editing the target logits before - sampling. The one-model path has no logits-editing hook, so each is + embedding_bias, top_p_decay and logits post-processor callbacks by editing + the target logits before sampling. The one-model path has no logits-editing hook, so each is rejected here instead of being silently dropped. Each check gates on a non-neutral value rather than on presence, since a @@ -252,6 +269,11 @@ def _validate_unsupported_logits_processors(request: LlmRequest) -> None: raise ValueError(UNSUPPORTED_NO_REPEAT_NGRAM_MSG) if getattr(request, "py_embedding_bias", None) is not None: raise ValueError(UNSUPPORTED_EMBEDDING_BIAS_MSG) + # The one-model sampler never calls a request's logits post-processors, and a fused + # verify (e.g. the MiniMax-M3 shard verify) does not even gather the full logits, so + # the callback would be skipped silently. + if getattr(request, "py_logits_post_processors", None): + raise ValueError(UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG) # Reuse the handler's own predicate so "active" cannot drift between paths. if top_p_decay_active(_request_get_sampling_params(request)): raise ValueError(UNSUPPORTED_TOP_P_DECAY_MSG) @@ -317,6 +339,68 @@ def __init__( new_tokens_lens=int_tensor((seq_slots,)), ) + def _device_slot_ids(self, seq_slots: list[int]) -> torch.Tensor: + """The slot ids on the device, uploaded only when they differ from the last upload. + + The buffer is this sampler's own: only this method writes it, so the last uploaded + list describes its contents, and the stream orders each upload after the scatter that + read the previous ids. + """ + n = len(seq_slots) + if self._slot_ids_cuda is None or self._slot_ids_cuda.numel() < n: + self._slot_ids_cuda = torch.empty(max(n, 8), dtype=torch.long, device="cuda") + self._slot_ids_host = None + if self._slot_ids_host != list(seq_slots): + self._slot_ids_cuda[:n].copy_( + torch.tensor(seq_slots, dtype=torch.long, pin_memory=prefer_pinned()), + non_blocking=True, + ) + self._slot_ids_host = list(seq_slots) + return self._slot_ids_cuda[:n] + + def _copy_outputs_to_host( + self, + o_verify_lens: Optional[torch.Tensor], + store_copies_on_side_stream: bool, + store_copies_done: Optional[torch.cuda.Event] = None, + ) -> tuple[SampleStateTensorsSpec, Optional[torch.cuda.Event]]: + """Host copies of the step's stored outputs, and the side-stream event to await them by (or None). + + With ``d2h_side_stream_copies`` the per-slot store is copied on the D2H side stream, so the copies + overlap the next step's input staging instead of preceding it on the main stream. That is sound + only while nothing rewrites the store before the sampler event is synchronized host-side, which + holds when this sampler stored the rows itself: the store is next written by the following + ``sample_async``, after the loop has synchronized this step. ``verify_lens`` is a forward output + buffer the next replay rewrites, so it is always copied on the main stream. + + ``store_copies_done`` (a step fold's ``store_copies_done``) is recorded on the side stream after the copies: + a forward that stores into the store waits on it first, which keeps the side-stream copies sound when the + next replay rather than the next ``sample_async`` writes the store. + """ + verify_lens = self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None + if not ( + self.d2h_side_stream_copies + and store_copies_on_side_stream + and not self._async_worker_active() + ): + host = SampleStateTensorsSpec( + new_tokens=self._copy_to_host(self.store.new_tokens), + new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), + next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), + verify_lens=verify_lens, + ) + return host, None + with self._make_side_stream_copier() as copier: + host = SampleStateTensorsSpec( + new_tokens=copier.stage_copy_to_host(self.store.new_tokens), + new_tokens_lens=copier.stage_copy_to_host(self.store.new_tokens_lens), + next_draft_tokens=copier.stage_copy_to_host(self.store.next_draft_tokens), + verify_lens=verify_lens, + ) + if store_copies_done is not None and copier.event is not None: + store_copies_done.record(self._d2h_side_stream) + return host, copier.event + def bind_spec_step_fold(self, spec_step_fold: SpecStepFold) -> None: """Let the model's forward store into this sampler's per-slot store (see SpecStepFold).""" self.spec_step_fold = spec_step_fold @@ -414,8 +498,11 @@ def _store_outputs( 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) + if self.static_slot_ids: + slots = self._device_slot_ids(seq_slots) + else: + 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 @@ -546,7 +633,10 @@ def sample_async( o_verify_lens = o_verify_lens[num_skip : num_skip + num_sampling_requests] step_fold = self.spec_step_fold - if step_fold is None or not step_fold.sampler_rows_stored(outputs, num_skip, seq_slots): + stored_here = step_fold is None or not step_fold.sampler_rows_stored( + outputs, num_skip, seq_slots + ) + if stored_here: self._store_outputs( seq_slots, o_new_tokens, o_next_new_tokens, o_new_tokens_lens, o_next_draft_tokens ) @@ -558,13 +648,15 @@ def sample_async( next_draft_tokens=self.store.next_draft_tokens, ) - host_tensors = SampleStateTensorsSpec( - new_tokens=self._copy_to_host(self.store.new_tokens), - new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), - next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), - verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None), + # A fold whose forward waits for the store copies (store_copies_done) lets them go on the side stream + # on the steps its forward stored the rows too. + store_copies_done = getattr(step_fold, "store_copies_done", None) if step_fold else None + host_tensors, side_stream_event = self._copy_outputs_to_host( + o_verify_lens, + store_copies_on_side_stream=stored_here or store_copies_done is not None, + store_copies_done=store_copies_done, ) - sampler_event = self._record_sampler_event() + sampler_event = self._record_sampler_event(side_stream_event) # Add dummy draft tokens to context requests for KV cache preparation for request in finished_context_requests: diff --git a/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py index a5c777b5cf7c..7c3ff30b3c30 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py @@ -18,11 +18,11 @@ criteria would quietly fall back to the built-in implementation, pass, and report the built-in's numbers as the target's. -The deployment is the one the targets route (``models/minimax_m3/routing.py``): the NVFP4 checkpoint on sm_103, -tensor-parallel attention and experts on 4 or 8 ranks, the MiniMax-M3 MSA sparse attention with the fused five-way -QKV/index projection, an FP8 indexer and an NVFP4 KV cache. Each target runs GSM8K without speculation and with -one-model Eagle3, where a chat-GSM8K acceptance probe also guards the drafter (accuracy alone does not notice a -corrupted draft KV). +The deployment is the one the targets route (``models/minimax_m3/routing.py``): the NVFP4 checkpoint on sm_103 +(GB300, 4 or 8 ranks) or sm_107 (Rubin, 2, 4 or 8 ranks), tensor-parallel attention and experts, the MiniMax-M3 MSA +sparse attention with the fused five-way QKV/index projection, an FP8 indexer and an NVFP4 KV cache. Each target runs +GSM8K without speculation and with one-model Eagle3, where a chat-GSM8K acceptance probe also guards the drafter +(accuracy alone does not notice a corrupted draft KV). """ import json @@ -48,7 +48,11 @@ # The targets assert their own SM at construction: certification is per GPU architecture. skip_not_sm103 = pytest.mark.skipif( get_sm_version() != 103, - reason="the MiniMax-M3 modeling_v2 targets are certified on sm_103 only", + reason="these MiniMax-M3 modeling_v2 targets are certified on sm_103 only", +) +skip_not_sm107 = pytest.mark.skipif( + get_sm_version() != 107, + reason="these MiniMax-M3 modeling_v2 targets are certified on sm_107 only", ) # The chat-GSM8K acceptance floors of the in-tree test on an NVFP4 KV cache (test_llm_api_pytorch.TestMiniMaxM3). @@ -166,3 +170,51 @@ def test_gsm8k(self): @pytest.mark.skip_less_mpi_world_size(8) def test_eagle3_gsm8k_and_acceptance(self): self._gsm8k(eagle3=True) + + +class TestModelingV2MiniMaxM3Nvfp4Sm107Tp2(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_107 / tp2: two Rubin GPUs.""" + + TP = 2 + + @skip_not_sm107 + @pytest.mark.skip_less_device(2) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm107 + @pytest.mark.skip_less_device(2) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) + + +class TestModelingV2MiniMaxM3Nvfp4Sm107Tp4(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_107 / tp4: four Rubin GPUs.""" + + TP = 4 + + @skip_not_sm107 + @pytest.mark.skip_less_device(4) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm107 + @pytest.mark.skip_less_device(4) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) + + +class TestModelingV2MiniMaxM3Nvfp4Sm107Tp8(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_107 / tp8: eight Rubin GPUs over two trays (an 8-rank MPI world across the two nodes).""" + + TP = 8 + + @skip_not_sm107 + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm107 + @pytest.mark.skip_less_mpi_world_size(8) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 1b29781b38ac..82cd90619c26 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -72,6 +72,12 @@ l0_cpu: - unittest/_torch/modeling_v2/test_modeling_v2_routing.py - unittest/_torch/modeling_v2/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py + # MiniMax-M3 sm_107 tp2 target: its device-free (cpu_only) tests. + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py + - unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_moe_reduces_only_local_terms diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py new file mode 100644 index 000000000000..fac388a85ec1 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 V2 cache manager: the SM107 per-token estimate adds what the generic sum leaves out.""" + +import types + +import pytest + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import cache_manager + +KVM = cache_manager.MiniMaxM3KVCacheManagerV2 + + +def _model_config(nvfp4=True, index_dtype="fp8"): + quant_mode = types.SimpleNamespace( + has_fp4_kv_cache=lambda: nvfp4, has_fp8_kv_cache=lambda: not nvfp4 + ) + return types.SimpleNamespace( + pretrained_config=types.SimpleNamespace( + num_hidden_layers=60, + num_key_value_heads=4, + head_dim=128, + hidden_size=6144, + num_attention_heads=64, + ), + sparse_attention_config=types.SimpleNamespace( + sparse_index_dim=128, indexer_kv_dtype=index_dtype + ), + quant_config=types.SimpleNamespace(quant_mode=quant_mode), + ) + + +def _mapping(tp_size, attention_dp): + return types.SimpleNamespace(tp_size=tp_size, pp_size=1, enable_attention_dp=attention_dp) + + +@pytest.fixture +def one_draft_layer(monkeypatch): + monkeypatch.setattr( + cache_manager, "shared_draft_layer_count", lambda spec_config, layer_mask: 1 + ) + + +@pytest.mark.parametrize( + "attention_dp,expected_extra", + [ + # TP2 attention, 2 KV heads per rank: index-K 57 x 128 + dense 3 x (512 - 288) + draft 512. + (False, 57 * 128 + 3 * (512 - 288) + 512), + # Attention DP, 4 KV heads per rank: index-K 57 x 128 + dense 3 x (1024 - 576) + draft 1024. + (True, 57 * 128 + 3 * (1024 - 576) + 1024), + ], +) +def test_extra_bytes_per_token_nvfp4(one_draft_layer, attention_dp, expected_extra): + extra = KVM._extra_bytes_per_token( + _model_config(), _mapping(2, attention_dp), None, {"spec_config": object()} + ) + assert extra == expected_extra + + +def test_separate_draft_cache_not_double_counted(one_draft_layer): + kwargs = {"spec_config": object(), "use_separate_draft_kv_cache": True} + extra = KVM._extra_bytes_per_token(_model_config(), _mapping(2, False), None, kwargs) + assert extra == 57 * 128 + 3 * (512 - 288) + + +def test_estimate_gate(monkeypatch): + calls = [] + + def base(cls, model_config, mapping, num_layers=None, **kwargs): + calls.append(kwargs) + return (17280, 0) + + monkeypatch.setattr( + cache_manager.KVCacheManagerV2, "get_cache_size_per_token", classmethod(base) + ) + monkeypatch.setattr(KVM, "_extra_bytes_per_token", classmethod(lambda cls, *a: 8480)) + monkeypatch.setattr(cache_manager, "get_sm_version", lambda: 107) + assert KVM.get_cache_size_per_token(None, None) == (25760, 0) + assert KVM.get_cache_size_per_token(None, None, is_draft=True) == (17280, 0) + monkeypatch.setattr(cache_manager, "get_sm_version", lambda: 103) + assert KVM.get_cache_size_per_token(None, None) == (17280, 0) + monkeypatch.setattr(cache_manager, "get_sm_version", lambda: 107) + monkeypatch.setattr(cache_manager, "_SM107_EXACT_KV_SIZE_ESTIMATE", False) + assert KVM.get_cache_size_per_token(None, None) == (17280, 0) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py index 88760a79452f..07376d86d2d7 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py @@ -1158,7 +1158,7 @@ def test_resolve_num_topk_chunks_is_shape_only_power_of_two(total_q, num_kv_head (33, {}), ], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_launch_options(local_batch, expected, capability): if capability != (10, 3): expected = {} @@ -1192,7 +1192,7 @@ def test_sm100f_nvfp4_launch_options(local_batch, expected, capability): (15, None), ], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_num_topk_chunks(local_batch, expected, capability): if capability != (10, 3): expected = None @@ -1213,7 +1213,7 @@ def test_sm100f_nvfp4_num_topk_chunks(local_batch, expected, capability): ("local_batch", "expected"), [(10, 1), (11, 2), (12, 2), (14, 2), (15, 1)], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_query_group_size_is_narrowly_scoped(local_batch, expected, capability): if capability != (10, 3): expected = 1 @@ -1259,7 +1259,7 @@ def test_sm100f_nvfp4_launch_options_rejects_unmeasured_shapes(): (8, 14, False), ], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_linear_softmax_policy(num_kv_heads, local_batch, expected, capability): if capability != (10, 3): expected = False @@ -1290,7 +1290,7 @@ def test_sm100f_nvfp4_linear_softmax_policy(num_kv_heads, local_batch, expected, (12, {}), ], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_merge_launch_options_is_narrowly_scoped(local_batch, expected, capability): if capability != (10, 3): expected = {} 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 05e053be775d..72e6ea2c2b5b 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -48,14 +48,18 @@ @pytest.mark.cpu_only +@pytest.mark.parametrize( + "tail_in_forward", [True, False], ids=["tail_in_forward", "tail_in_prepare"] +) @pytest.mark.parametrize("empty_state", ["empty_batch", "no_requests", "no_lengths"]) def test_msa_metadata_clears_padded_cache_slot_tail( - monkeypatch: pytest.MonkeyPatch, empty_state: str + monkeypatch: pytest.MonkeyPatch, empty_state: str, tail_in_forward: bool ) -> None: """A smaller replay must not reuse the previous step's live cache slots.""" from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend monkeypatch.setattr(msa_backend, "maybe_pin_memory", lambda tensor: tensor) + monkeypatch.setattr(msa_backend, "TAIL_FILL_IN_FORWARD", tail_in_forward) metadata_cls = MiniMaxM3MsaSparseAttention.Metadata metadata = metadata_cls.__new__(metadata_cls) metadata._msa_buffers_ready = True @@ -81,6 +85,8 @@ def test_msa_metadata_clears_padded_cache_slot_tail( # live count off it. metadata._num_tokens = count metadata._build_msa_fields() + # TAIL_FILL_IN_FORWARD: the forward (on_update_kv_lens) fills the tail before anything reads it. + metadata._fill_tail_in_forward() assert metadata.msa_out_cache_loc.tolist() == list(range(12, 12 + count)) + [-1] * ( 4 - count ) @@ -259,7 +265,7 @@ def test_nvfp4_sparse_decode_rejects_unsupported_device( block_scale = torch.empty((1, 1, 128, 8), dtype=torch.uint8) global_scale = torch.ones(1) - with pytest.raises(NotImplementedError, match="NVFP4 sparse decode requires SM100/SM103"): + with pytest.raises(NotImplementedError, match="NVFP4 sparse decode requires SM100/SM103/SM107"): minimax_m3_sparse_attn_decode( q, kv, @@ -278,6 +284,34 @@ def test_nvfp4_sparse_decode_rejects_unsupported_device( get_capability.assert_called_once_with(q.device) +@pytest.mark.cpu_only +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) +def test_nvfp4_sparse_decode_admits_the_sm100_family( + monkeypatch: pytest.MonkeyPatch, capability: tuple[int, int] +) -> None: + """SM100-family devices, Rubin included, pass the device check and reach + input validation, which here rejects the missing scales.""" + get_capability = Mock(spec_set=torch.cuda.get_device_capability, return_value=capability) + monkeypatch.setattr(torch.cuda, "get_device_capability", get_capability) + q = torch.empty((1, 8, 128), dtype=torch.bfloat16) + kv = torch.empty((1, 1, 128, 64), dtype=torch.uint8) + + with pytest.raises(ValueError, match="needs all four"): + minimax_m3_sparse_attn_decode( + q, + kv, + kv, + torch.zeros((1, 1, 1), dtype=torch.int32), + torch.zeros((1, 1), dtype=torch.int32), + torch.ones(1, dtype=torch.int32), + sm_scale=128**-0.5, + output=torch.empty_like(q), + decode_query_len=1, + k_block_scale=torch.empty((1, 1, 128, 8), dtype=torch.uint8), + ) + get_capability.assert_called_once_with(q.device) + + @pytest.mark.cpu_only @pytest.mark.parametrize("indexer_kv_dtype", ["bf16", "fp8"]) @pytest.mark.parametrize("implementation", ["msa", "triton"]) @@ -490,6 +524,45 @@ def test_msa_buffers_stage_local_cache_views(monkeypatch: pytest.MonkeyPatch, nv torch.testing.assert_close(index_cache, torch.full_like(index_cache, 3)) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("kv_dtype", "expect_lengths"), + [ + (DataType.FP8, False), + (DataType.NVFP4, True), + ], +) +def test_msa_buffers_reserve_cumulative_lengths_for_csr_prefill( + monkeypatch: pytest.MonkeyPatch, + kv_dtype: DataType, + expect_lengths: bool, +) -> None: + """The NVFP4 CSR kernel reads graph-stable cumulative lengths, so only an + NVFP4 cache reserves them.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend + + metadata = _buffer_metadata(dtype=kv_dtype, sparse_layer_ids=[], layer_offsets={}) + reserved = {} + + def get_empty(buffers, shape, **kwargs): + reserved[kwargs["cache_name"]] = tuple(shape) + return torch.empty(shape, dtype=kwargs["dtype"]) + + monkeypatch.setattr(metadata, "get_empty", get_empty) + monkeypatch.setattr(msa_backend, "uniform_subpages_per_slot", lambda manager: 0) + metadata._create_msa_buffers() + + names = ("msa_qo_lens_cuda", "msa_cu_q_lens", "msa_cu_kv_lens") + if expect_lengths: + assert reserved["msa_qo_lens_cuda"] == (MAX_NUM_SEQUENCES,) + assert reserved["msa_cu_q_lens"] == (MAX_NUM_SEQUENCES + 1,) + assert reserved["msa_cu_kv_lens"] == (MAX_NUM_SEQUENCES + 1,) + assert all(getattr(metadata, name) is not None for name in names) + else: + assert not set(names) & set(reserved) + assert all(getattr(metadata, name) is None for name in names) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( ("factors", "expected"), diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py index 2f63cd962e39..d0e6d45179aa 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py @@ -9,16 +9,23 @@ from unittest.mock import create_autospec +import numpy as np import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3KVCacheManagerV2 +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + msa_backend, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.common import ( MiniMaxM3SparseMetadataParams, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( MSA_REQUIRED_TOPK, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( + write_subpage_block_table, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( MiniMaxM3MsaSparseAttentionMetadata, _fill_unused_slots, @@ -30,6 +37,7 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -48,8 +56,8 @@ def fake_base_init(self, *args, **kwargs) -> None: self.is_disagg = False self.dtype = DataType.BF16 self.layer_offsets = {3: 0} - # _prepare_page_table_tensor's result for the two-pool host table the copy tests build. - self.num_attention_op_pools = 2 + # Set by the base's page-table setup, which this fake base skips. + self.num_pools = self.num_attention_op_pools = 1 def no_index_k_buffer(self, layer_idx, **kwargs): del self, layer_idx, kwargs @@ -58,10 +66,10 @@ def no_index_k_buffer(self, layer_idx, **kwargs): monkeypatch.setattr(KVCacheManagerV2, "__init__", fake_base_init) monkeypatch.setattr(KVCacheManagerV2, "get_index_k_buffer", no_index_k_buffer) monkeypatch.setattr(MiniMaxM3KVCacheManagerV2, "_compute_num_total_slots", lambda self: 0) - # The constructor resolves every layer's K/V view, which needs a real pool. Stub it there - # only: the pool-view tests below call get_buffers against their own geometry. - with monkeypatch.context() as construction: - construction.setattr(MiniMaxM3KVCacheManagerV2, "get_buffers", lambda self, *args: None) + # The constructor resolves each layer's K/V views; the bare base has no pools to resolve them + # from. Stub that only while constructing, so the pool-view tests below build real views. + with monkeypatch.context() as construct: + construct.setattr(MiniMaxM3KVCacheManagerV2, "get_buffers", lambda self, *args: None) return MiniMaxM3KVCacheManagerV2( num_layers=4, sparse_attention_config=MiniMaxM3SparseAttentionConfig() ) @@ -373,3 +381,203 @@ def test_unused_slots_are_filled_after_a_full_buffer_step(): buf[:3] = 0 _fill_unused_slots(buf, 3) assert buf.tolist() == [0, 0, 0, -1, -1, -1, -1, -1] + + +class _PageIndexView: + """A request's KV cache as get_block_ids_per_seq reads it: an int32 page-index view per pool.""" + + def __init__(self, view: np.ndarray): + self._view = view + + def get_base_page_indices(self, pool_id: int) -> np.ndarray: + assert pool_id == 0 + return self._view + + +def _page_index_views(seed: int): + rng = np.random.default_rng(seed) + batch = int(rng.integers(0, 5)) + same_width = rng.random() < 0.5 + width = int(rng.integers(1, 9000)) + views = {} + for rid in range(batch): + w = width if same_width else int(rng.integers(0, 9000)) + view = rng.integers(0, 2**31 - 1, w, dtype=np.int64).astype(np.int32) + view[int(rng.integers(0, w + 1)) :] = BAD_PAGE_INDEX + view[rng.random(w) < float(rng.choice([0.0, 0.01, 0.5]))] = BAD_PAGE_INDEX + view.setflags(write=False) + views[100 + rid] = _PageIndexView(view) + return views + + +def _bare_manager() -> MiniMaxM3KVCacheManagerV2: + """A manager with only the state get_block_ids_per_seq reads (no pools, no V2 base).""" + return object.__new__(MiniMaxM3KVCacheManagerV2) + + +def test_block_ids_from_views_is_off_by_default(): + assert not _bare_manager().block_ids_from_views + + +@pytest.mark.parametrize("seed", range(24)) +def test_block_ids_from_views_matches_the_list_build(seed): + manager = _bare_manager() + manager.kv_cache_map = _page_index_views(seed) + request_ids = list(manager.kv_cache_map) + + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=False, reuse_pool_views=False) + built_from_lists = manager.get_block_ids_per_seq(request_ids) + manager.configure_decode_shortcuts( + skip_unchanged_block_offsets=False, reuse_pool_views=False, block_ids_from_views=True + ) + copied_from_views = manager.get_block_ids_per_seq(request_ids) + + assert copied_from_views.dtype == built_from_lists.dtype == torch.int32 + assert copied_from_views.is_pinned() == built_from_lists.is_pinned() + assert torch.equal(copied_from_views, built_from_lists) + + +# --------------------------------------------------------------------------- +# Sub-page block table expanded inside the forward +# --------------------------------------------------------------------------- + + +def _subpage_metadata( + rows: int, max_blocks: int, factor: int +) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + metadata._msa_subpage_rows = rows # as prepare() stages them + metadata.msa_block_table = torch.randint( + 0, 1 << 20, (rows + 2, max_blocks), dtype=torch.int32, device="cuda" + ) + metadata.msa_subpage_block_table = torch.full( + (rows + 2, 2, max_blocks), -7, dtype=torch.int32, device="cuda" + ) + metadata._msa_subpages_per_slot = factor + metadata._msa_subpage_expanded = False + return metadata + + +def _expected_subpages(metadata, rows: int) -> torch.Tensor: + out = torch.empty_like(metadata.msa_subpage_block_table[:rows]) + write_subpage_block_table(metadata.msa_block_table[:rows], metadata._msa_subpages_per_slot, out) + return out + + +@requires_cuda +def test_the_first_reader_of_a_forward_expands_the_subpage_table(monkeypatch): + monkeypatch.setattr(msa_backend, "SUBPAGE_IN_GRAPH", True) + metadata = _subpage_metadata(rows=3, max_blocks=40, factor=2) + calls = [] + real = msa_backend.write_subpage_block_table + monkeypatch.setattr( + msa_backend, "write_subpage_block_table", lambda *a: (calls.append(1), real(*a)) + ) + + table, factor = metadata.msa_subpage_rows(0, 3) + again, _ = metadata.msa_subpage_rows(1, 3) + + assert factor == 2 and len(calls) == 1 + assert torch.equal(table, _expected_subpages(metadata, 3)) + assert torch.equal(again, table[1:3]) + # Rows past the step's requests are left alone. + assert metadata.msa_subpage_block_table[3:].eq(-7).all() + + +@requires_cuda +def test_each_forward_expands_the_current_block_table(monkeypatch): + monkeypatch.setattr(msa_backend, "SUBPAGE_IN_GRAPH", True) + metadata = _subpage_metadata(rows=2, max_blocks=16, factor=4) + metadata.msa_subpage_rows(0, 2) + metadata.msa_block_table[:2, :5] += 11 + # The next forward (on_update_kv_lens opens it). + metadata._msa_subpage_expanded = False + + table, _ = metadata.msa_subpage_rows(0, 2) + + assert torch.equal(table, _expected_subpages(metadata, 2)) + + +@requires_cuda +def test_a_captured_forward_expands_the_table_at_every_replay(monkeypatch): + monkeypatch.setattr(msa_backend, "SUBPAGE_IN_GRAPH", True) + metadata = _subpage_metadata(rows=4, max_blocks=32, factor=2) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + metadata._msa_subpage_expanded = False + metadata.msa_subpage_rows(0, 4) + metadata.msa_subpage_rows(2, 4) + torch.cuda.current_stream().wait_stream(stream) + + for step in range(3): + # The pack lands this step's block table; then the replay. + metadata.msa_block_table[:4].copy_( + torch.randint(0, 1 << 20, (4, 32), dtype=torch.int32, device="cuda") + ) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(metadata.msa_subpage_block_table[:4], _expected_subpages(metadata, 4)), ( + step + ) + + +def test_nothing_is_expanded_without_a_prepared_step(): + metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + metadata.msa_subpage_block_table = torch.arange(12, dtype=torch.int32).reshape(2, 2, 3) + metadata._msa_subpages_per_slot = 2 + + table, factor = metadata.msa_subpage_rows(0, 2) + + assert factor == 2 and torch.equal(table, torch.arange(12, dtype=torch.int32).reshape(2, 2, 3)) + + +# --------------------------------------------------------------------------- +# msa_out_cache_loc tail filled inside the forward +# --------------------------------------------------------------------------- + + +def _tail_metadata(cap: int, start) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + metadata.msa_out_cache_loc = torch.arange(100, 100 + cap, dtype=torch.int32, device="cuda") + metadata._msa_tail_start = start + return metadata + + +@requires_cuda +def test_the_forward_fills_only_past_the_live_slots(): + metadata = _tail_metadata(10, 4) + metadata._fill_tail_in_forward() + assert metadata.msa_out_cache_loc.tolist() == [100, 101, 102, 103] + [-1] * 6 + + untouched = _tail_metadata(6, None) + untouched._fill_tail_in_forward() + assert untouched.msa_out_cache_loc.tolist() == list(range(100, 106)) + + full = _tail_metadata(6, 6) + full._fill_tail_in_forward() + assert full.msa_out_cache_loc.tolist() == list(range(100, 106)) + + +@requires_cuda +def test_a_captured_forward_refills_its_tail_at_every_replay(): + """Another graph key sharing the pool block writes live slots past this key's tail between replays; each replay of + this key's forward must invalidate them again before its own readers run.""" + metadata = _tail_metadata(12, 5) + buf = metadata.msa_out_cache_loc + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + metadata._fill_tail_in_forward() + torch.cuda.current_stream().wait_stream(stream) + + for step in range(3): + buf[:5] = ( + torch.arange(5, dtype=torch.int32, device="cuda") + 10 * step + ) # this step's live slots + buf[5:9] = 777 # another key's in-graph writes into the shared block + graph.replay() + torch.cuda.synchronize() + assert buf.tolist() == [10 * step + i for i in range(5)] + [-1] * 7, step diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_skip_cache_routing.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_skip_cache_routing.py new file mode 100644 index 000000000000..9e50093d41ff --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_skip_cache_routing.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""No buffer a captured graph writes on the device is routed through the H2D skip cache (K3 audit #150, check #6). + +The skip cache (``pyexecutor/h2d_cache.py``) skips an upload when its host shadow says the destination already holds +the values. That is sound only while nothing else writes the destination: a graph replay that rewrites it on the device +(``on_update_kv_lens`` / ``msa.kv_update`` for the MSA lengths, slots and valid-block counts, ``_preprocess_inputs`` for +``kv_lens_cuda``) leaves the shadow describing values the device no longer holds. This test scans +``tensorrt_llm/_torch`` and fails if any of those buffers is the destination of a skipping helper; they have to go +through ``step_inputs.copy_`` or the never-skipping ``copy_h2d_pool_shared`` / ``fill_tail_pool_shared``. + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +# Buffers a captured forward writes on the device. +_DEVICE_WRITTEN = { + "msa_seq_lens_cuda", + "msa_cu_kv_lens", + "msa_out_cache_loc", + "msa_n_valid_blocks", + "_msa_prefill_n_valid_blocks", + "kv_lens_cuda", +} +# Helpers that may skip the write; their first argument is the destination. +_SKIPPING = {"copy_h2d_if_changed", "fill_tail_if_changed"} + + +def _receiver_names(node) -> set: + """Attribute / variable names along the receiver chain of an expression (``self.a[:n].view(...)`` -> {"a"}).""" + names = set() + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute): + names.add(node.attr) + node = node.func if isinstance(node, ast.Call) else node.value + if isinstance(node, ast.Name): + names.add(node.id) + return names + + +def _skipped_device_buffers(tree: ast.AST): + hits = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not node.args: + continue + f = node.func + name = f.id if isinstance(f, ast.Name) else (f.attr if isinstance(f, ast.Attribute) else "") + if name in _SKIPPING: + bad = _receiver_names(node.args[0]) & _DEVICE_WRITTEN + if bad: + hits.append((node.lineno, name, sorted(bad))) + return hits + + +def test_no_device_written_buffer_goes_through_the_skip_cache() -> None: + found = [] + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if not any(helper in src for helper in _SKIPPING): + continue + for line, helper, bad in _skipped_device_buffers(ast.parse(src)): + found.append(f"{path.relative_to(_ROOT).as_posix()}:{line} {helper}({', '.join(bad)})") + assert not found, "device-written buffers routed through the H2D skip cache:\n" + "\n".join( + found + ) + + +def test_the_scan_sees_a_skipped_device_buffer() -> None: + """The scan itself: a skipping helper writing a device-written buffer is reported, a never-skipping one is not.""" + bad = ast.parse( + "copy_h2d_if_changed(self.msa_seq_lens_cuda[:b], host)\nfill_tail_if_changed(m.msa_out_cache_loc, n, -1)" + ) + ok = ast.parse( + "copy_h2d_pool_shared(self.msa_seq_lens_cuda[:b], host)\ncopy_h2d_if_changed(self.msa_block_table[:b], h)" + ) + assert [h[1] for h in _skipped_device_buffers(bad)] == [ + "copy_h2d_if_changed", + "fill_tail_if_changed", + ] + assert _skipped_device_buffers(ok) == [] diff --git a/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py b/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py index a9e1c62b1e0e..83ab912fdad6 100644 --- a/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py +++ b/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py @@ -35,6 +35,10 @@ # A head width the SM100-family kernel table does carry for 16-bit paged-context # attention. _SUPPORTED_HEAD_SIZE = 64 +# The page size and head width of MiniMax-M3's dense and draft attention. The +# kernel table carries no 128-token-page kernels; SM107 generates them at runtime. +_NVRTC_TOKENS_PER_BLOCK = 128 +_NVRTC_HEAD_SIZE = 128 _NO_GPU = not torch.cuda.is_available() _NO_GPU_REASON = "needs a GPU; missing bindings on a GPU are a test failure" @@ -188,3 +192,30 @@ def test_lookup_is_page_size_sensitive(): ) assert present assert not odd_page + + +@pytest.mark.parametrize("kv_cache_dtype", [DataType.BF16, DataType.FP8], ids=["bf16", "fp8"]) +def test_lookup_reports_present_for_sm107_128_token_pages(kv_cache_dtype): + """SM107 serves dense paged context over 128-token pages with BF16 output. + + No precompiled kernel exists for that page size, so matched BF16/E4M3 Q + and KV with BF16 output must be generated at runtime rather than reported + absent, which would refuse paged-context attention for such models. The + runtime path is limited to that output precision: FP8 output at the same + page size has no kernel either way and must still report absent. + """ + if get_sm_version() != 107: + pytest.skip("runtime-generated 128-token-page context kernels are specific to SM107") + assert thop.fused_context_fmha_kernel_exists( + head_size=_NVRTC_HEAD_SIZE, + kv_cache_dtype=kv_cache_dtype, + tokens_per_block=_NVRTC_TOKENS_PER_BLOCK, + output_dtype=DataType.BF16, + ) + if kv_cache_dtype == DataType.FP8: + assert not thop.fused_context_fmha_kernel_exists( + head_size=_NVRTC_HEAD_SIZE, + kv_cache_dtype=kv_cache_dtype, + tokens_per_block=_NVRTC_TOKENS_PER_BLOCK, + output_dtype=DataType.FP8, + ) diff --git a/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py new file mode 100644 index 000000000000..479d5f279cfa --- /dev/null +++ b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FallbackFmha's refusal of paged prefill at a KV manager's TRTLLM-Gen extra page sizes: kept on every SM but SM107, +where no TRTLLM-Gen FMHA exists and the thop paged-context kernel serves those pages (its SM107 P128 oracle is +test_modeling_v2_m3_draft_attention_paged_context.py). CPU only: the check runs before any kernel.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.attention.backends.fmha import fallback +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase + +pytestmark = pytest.mark.cpu_only + + +class _PastRefusal(Exception): + """Raised by the stubbed paged-context validation that follows the refusal check.""" + + +def _fmha(): + fmha = object.__new__(fallback.FallbackFmha) + + def _validate(metadata): + raise _PastRefusal() + + fmha._validate_paged_context_fmha = _validate + return fmha + + +def _context_batch(tokens_per_block: int = 128, paged: bool = True): + return SimpleNamespace( + num_contexts=1, + kv_cache_manager=SimpleNamespace(trtllm_gen_extra_tokens_per_block=frozenset({128})), + use_paged_context_fmha=paged, + tokens_per_block=tokens_per_block, + ) + + +@pytest.mark.parametrize("sm", [90, 100, 103, 120]) +def test_extra_page_paged_prefill_is_refused_off_sm107(monkeypatch, sm) -> None: + monkeypatch.setattr(fallback, "get_sm_version", lambda: sm) + with pytest.raises(RuntimeError, match="cannot serve paged prefill with 128-token KV pages"): + _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT) + + +def test_extra_page_paged_prefill_is_served_on_sm107(monkeypatch) -> None: + monkeypatch.setattr(fallback, "get_sm_version", lambda: 107) + with pytest.raises(_PastRefusal): + _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT) + + +@pytest.mark.parametrize("sm", [100, 107]) +def test_other_page_sizes_and_unpaged_prefill_pass_the_check(monkeypatch, sm) -> None: + monkeypatch.setattr(fallback, "get_sm_version", lambda: sm) + for batch in (_context_batch(tokens_per_block=32), _context_batch(paged=False)): + with pytest.raises(_PastRefusal): + _fmha()._is_supported(None, None, None, batch, None, phase=FmhaPhase.CONTEXT) diff --git a/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py b/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py new file mode 100644 index 000000000000..319bb4abd8b6 --- /dev/null +++ b/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TRTLLM_AR_TUNE_NCCL_SYMMETRIC: the AllReduce tuner's candidate list with and without NCCL_SYMMETRIC.""" + +import pytest +import torch + +from tensorrt_llm._torch.autotuner import OptimizationProfile +from tensorrt_llm._torch.custom_ops import torch_custom_ops +from tensorrt_llm.functional import AllReduceStrategy + + +def _runner(): + runner = object.__new__(torch_custom_ops.AllReduceRunner) + runner.tp_size = 2 + return runner + + +@pytest.mark.parametrize("rows", [8, 4096, 16384]) +@pytest.mark.parametrize("tune_symmetric", [True, False]) +def test_nccl_symmetric_candidate_follows_knob( + monkeypatch, rows: int, tune_symmetric: bool +) -> None: + monkeypatch.setattr(torch_custom_ops, "_AR_TUNE_NCCL_SYMMETRIC", tune_symmetric) + inputs = [torch.empty(rows, 6144, dtype=torch.bfloat16)] + tactics = _runner().get_valid_tactics(inputs, OptimizationProfile()) + assert AllReduceStrategy.NCCL.value in tactics + assert (AllReduceStrategy.NCCL_SYMMETRIC.value in tactics) is tune_symmetric + # Off removes only NCCL_SYMMETRIC: the other candidates are unchanged. + monkeypatch.setattr(torch_custom_ops, "_AR_TUNE_NCCL_SYMMETRIC", True) + stock = _runner().get_valid_tactics(inputs, OptimizationProfile()) + assert [t for t in stock if t != AllReduceStrategy.NCCL_SYMMETRIC.value] == [ + t for t in tactics if t != AllReduceStrategy.NCCL_SYMMETRIC.value + ] diff --git a/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py index 9fda60feb590..13574cd0117e 100644 --- a/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py +++ b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """DecoderRunner and DecoderMetadata wiring of the step-loop hooks a model declares.""" +from types import SimpleNamespace + import pytest from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( @@ -30,12 +32,16 @@ def _runner( - sparse_attention_config: BaseSparseAttentionConfig | None, hooks: RuntimeHooks + sparse_attention_config: BaseSparseAttentionConfig | None, + hooks: RuntimeHooks, + max_draft_len: int = 0, ) -> DecoderRunner: - # Only the two attributes the threshold reads. + # Only the attributes the threshold reads. runner = object.__new__(DecoderRunner) runner.sparse_attention_config = sparse_attention_config runner.runtime_hooks = hooks + runner._config = SimpleNamespace(max_draft_len=max_draft_len) + runner.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 return runner @@ -64,10 +70,14 @@ def test_a_sparse_config_that_asks_for_the_split_sets_the_threshold(): assert _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() == 512 -def test_short_identity_threshold_is_the_topk_span(): +@pytest.mark.parametrize("max_draft_len", [0, 3]) +def test_short_identity_threshold_keeps_one_step_of_headroom_in_the_topk_span(max_draft_len): + """The short graph family runs the identity selection: one generation step's tokens below the top-k span, the most + the graph key's host length estimate can lag the attended length under the overlap scheduler.""" config = MiniMaxM3SparseAttentionConfig(sparse_topk_blocks=16, sparse_block_size=128) - assert _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() == 16 * 128 + threshold = _runner(config, _SHORT_IDENTITY, max_draft_len)._short_long_seq_len_threshold() + assert threshold == 16 * 128 - (max_draft_len + 1) @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/executor/test_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_cuda_graph_runner.py index b0934702e4f9..b8b96b642ac6 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_runner.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_runner.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -16,6 +17,7 @@ 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._torch.speculative.spec_sampler_base import SampleStateTensorsSpec from tensorrt_llm.bindings import SamplingConfig from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, DeepSeekSparseAttentionConfig from tensorrt_llm.mapping import Mapping @@ -274,3 +276,46 @@ 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 + + +def _overlap_spec_batch(max_draft_len: int, *num_tokens: int) -> tuple: + """A decode batch under the overlap scheduler with speculative decoding: each request has a previous batch, and the + in-flight step's sampled tokens and drafts are device tensors the host has not read yet.""" + batch = _generation_batch(*num_tokens) + for request in batch.generation_requests: + request.py_batch_idx = 0 + request.py_is_first_draft = False + tensors = object.__new__(SampleStateTensorsSpec) + tensors.new_tokens = torch.zeros(1, len(num_tokens), 1, dtype=torch.int32) + tensors.next_draft_tokens = torch.zeros(len(num_tokens), max_draft_len, dtype=torch.int32) + return batch, tensors + + +@pytest.mark.parametrize("max_draft_len", [1, 3]) +def test_identity_short_mode_keeps_every_attended_token_within_the_topk_span(max_draft_len) -> None: + """With the identity threshold (top-k span minus one generation step's tokens), a decode step the runner keys short + attends at most top-k blocks for every verify token, whatever the in-flight step accepted. The host knows a + request's tokens up to the step before the in-flight one; the in-flight step accepts a in [1, draft + 1] tokens and + this step adds draft more, so its last verify token attends T + a + draft tokens. The bare top-k span let a + request attend one block past it (bug B: identity steps at 2049 tokens).""" + span = 16 * 128 + tokens_per_step = max_draft_len + 1 + spec_config = SimpleNamespace(max_draft_len=max_draft_len, _use_shared_kv_cache=True) + + def worst_attended(host_tokens: int) -> int: + return host_tokens + tokens_per_step + max_draft_len + + def short(threshold: int, host_tokens: int) -> bool: + runner = _decoder_runner(short_long_seq_len_threshold=threshold) + runner.spec_config = spec_config + batch, tensors = _overlap_spec_batch(max_draft_len, host_tokens) + return runner._get_seq_len_mode(batch, tensors) + + for host_tokens in range(span - 4 * tokens_per_step, span + 1): + if short(span - tokens_per_step, host_tokens): + assert worst_attended(host_tokens) <= span, host_tokens + # The unpadded span keys a request short whose worst case attends past it. + assert any( + short(span, t) and worst_attended(t) > span + for t in range(span - 4 * tokens_per_step, span + 1) + ) diff --git a/tests/unittest/_torch/executor/test_first_content_emission.py b/tests/unittest/_torch/executor/test_first_content_emission.py new file mode 100644 index 000000000000..932a5c130b1f --- /dev/null +++ b/tests/unittest/_torch/executor/test_first_content_emission.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the streaming emission schedule in ``PyExecutor``. + +``_should_emit_response`` decides on which decoding iterations a streaming +request gets a response. With ``stream_interval = S`` it emits on iteration 1, +on finish and every S iterations; it also emits on iterations +``2..first_content_emit_iters`` (the model's ``RuntimeHooks.first_content_emit_iters``; 1, the +default, keeps the stock schedule; the MiniMax-M3 sm_107 targets set 2), +so text held back by a reasoning parser after the first token reaches the +client one step later instead of one stream interval later. These tests call +the real method unbound against a minimal fake ``self`` and check that: + +* with ``first_content_emit_iters = 1`` (the default) the schedule is the stock + one: iteration 1, every S iterations, and finish; +* ``first_content_emit_iters = 2`` adds iteration 2 only; +* larger values add iterations ``2..N``; +* a finished request always emits, and ``emit_first_iter=False`` (early + first-token emission) still suppresses only a non-terminal iteration 1; +* ``stream_interval = 1`` emits every iteration whatever the setting. +""" + +from __future__ import annotations + +import types + +import pytest + +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks, runtime_hooks_of + +pytestmark = pytest.mark.cpu_only + + +def _fake_self(stream_interval: int, first_content_emit_iters: int): + return types.SimpleNamespace( + stream_interval=stream_interval, first_content_emit_iters=first_content_emit_iters + ) + + +def _request(decoding_iter: int, finished: bool = False): + return types.SimpleNamespace(py_decoding_iter=decoding_iter, is_finished=finished) + + +def _emitted_iters(stream_interval, first_content_emit_iters, last=60, emit_first_iter=True): + fake = _fake_self(stream_interval, first_content_emit_iters) + return [ + it + for it in range(1, last + 1) + if PyExecutor._should_emit_response(fake, _request(it), emit_first_iter) + ] + + +def _stock_iters(stream_interval, last=60): + return [it for it in range(1, last + 1) if it == 1 or it % stream_interval == 0] + + +@pytest.mark.parametrize("stream_interval", [2, 3, 20]) +def test_kill_matches_stock_schedule(stream_interval): + assert _emitted_iters(stream_interval, 1) == _stock_iters(stream_interval) + + +@pytest.mark.parametrize("stream_interval", [3, 20]) +def test_default_adds_iteration_two_only(stream_interval): + expected = sorted(set(_stock_iters(stream_interval)) | {2}) + assert _emitted_iters(stream_interval, 2) == expected + + +def test_larger_window_adds_iterations_up_to_n(): + expected = sorted(set(_stock_iters(20)) | {2, 3, 4, 5}) + assert _emitted_iters(20, 5) == expected + + +@pytest.mark.parametrize("first_content_emit_iters", [1, 2, 5]) +@pytest.mark.parametrize("decoding_iter", [1, 2, 7, 19]) +def test_finished_request_always_emits(first_content_emit_iters, decoding_iter): + fake = _fake_self(20, first_content_emit_iters) + for emit_first_iter in (True, False): + assert PyExecutor._should_emit_response( + fake, _request(decoding_iter, finished=True), emit_first_iter + ) + + +@pytest.mark.parametrize("first_content_emit_iters", [1, 2, 5]) +def test_early_first_token_suppresses_only_iteration_one(first_content_emit_iters): + with_first = _emitted_iters(20, first_content_emit_iters) + without_first = _emitted_iters(20, first_content_emit_iters, emit_first_iter=False) + assert without_first == [it for it in with_first if it != 1] + + +@pytest.mark.parametrize("first_content_emit_iters", [1, 2, 5]) +def test_stream_interval_one_emits_every_iteration(first_content_emit_iters): + assert _emitted_iters(1, first_content_emit_iters) == list(range(1, 61)) + + +class _PlainModel: + """A model that declares no runtime hooks (every model but the MiniMax-M3 sm_107 targets).""" + + +class _M3Model: + """A model whose runtime hooks opt into the iteration-2 emission, as the MiniMax-M3 sm_107 targets do.""" + + runtime_hooks = RuntimeHooks(first_content_emit_iters=2) + + +def test_runtime_hooks_choose_the_schedule(): + """PyExecutor reads the schedule from the model's runtime hooks: a model without hooks keeps the stock schedule + (1, 20, 40, ...), one that opts in also emits on iteration 2.""" + plain = runtime_hooks_of(_PlainModel()).first_content_emit_iters + m3 = runtime_hooks_of(_M3Model()).first_content_emit_iters + + assert _emitted_iters(20, plain) == [1, 20, 40, 60] + assert _emitted_iters(20, m3) == [1, 2, 20, 40, 60] diff --git a/tests/unittest/_torch/executor/test_h2d_cache.py b/tests/unittest/_torch/executor/test_h2d_cache.py index d70f58315ff2..3ea090968718 100644 --- a/tests/unittest/_torch/executor/test_h2d_cache.py +++ b/tests/unittest/_torch/executor/test_h2d_cache.py @@ -7,11 +7,14 @@ import pytest import torch +from tensorrt_llm._torch.memory_buffer_utils import Buffers from tensorrt_llm._torch.pyexecutor import h2d_cache from tensorrt_llm._torch.pyexecutor.h2d_cache import ( H2DSkipCache, copy_h2d_if_changed, + copy_h2d_pool_shared, fill_tail_if_changed, + fill_tail_pool_shared, ) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -314,3 +317,102 @@ def test_a_tail_record_keeps_its_storage_so_a_new_buffer_cannot_inherit_it(cache assert fresh.data_ptr() != address fill_tail_if_changed(fresh, 3, -1) assert fresh[3:].eq(-1).all() + + +# --------------------------------------------------------------------------- +# CUDA graph metadata pool: one block shared by every graph key +# --------------------------------------------------------------------------- + +NUM_SLOTS = 8 + + +def _graph_key_views(name): + """Two graph keys' metadata views from one pool, as create_cuda_graph_metadata hands them out.""" + pool = Buffers() + large = pool.get_buffer([NUM_SLOTS], torch.int32, name, reserve_buffer=True) + small = pool.get_buffer([NUM_SLOTS], torch.int32, name, reserve_buffer=True) + assert small.data_ptr() == large.data_ptr() + return large, small + + +def _capture_live_row_write(buf, rows, values): + """A graph whose replay writes buf[:rows] on the device, as on_update_kv_lens writes the live slots.""" + slots = torch.tensor(values, dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + buf[:rows].copy_(slots) + torch.cuda.current_stream().wait_stream(stream) + return graph + + +def _key_b_step(buf, live): + """Key B's prepare: its live slots, then -1 over the padded rows its captured producers also run.""" + copy_h2d_pool_shared(buf[: len(live)], _host(live)) + fill_tail_pool_shared(buf, len(live), -1) + + +@pytest.mark.parametrize("skip_pool_shared", [False, True], ids=["fixed", "kill_switch"]) +def test_two_graph_keys_over_one_block_keep_the_padded_rows_unused( + cache, monkeypatch, skip_pool_shared +): + monkeypatch.setattr(h2d_cache, "SKIP_POOL_SHARED_UPLOADS", skip_pool_shared) + key_a, key_b = _graph_key_views("msa_out_cache_loc") + replay_a = _capture_live_row_write(key_a, 6, [10, 11, 12, 13, 14, 15]) + + _key_b_step(key_b, [20, 21, 22]) + torch.cuda.synchronize() + assert key_b.tolist() == [20, 21, 22, -1, -1, -1, -1, -1] + # Key A's step: its replay writes six live slots into the shared block. + replay_a.replay() + # Key B's next step, with the same live slots. + _key_b_step(key_b, [20, 21, 22]) + torch.cuda.synchronize() + if skip_pool_shared: + # The old behaviour: both writes were skipped, so key B's live and padded rows hold key A's live slots. + assert key_b.tolist() == [10, 11, 12, 13, 14, 15, -1, -1] + else: + assert key_b.tolist() == [20, 21, 22, -1, -1, -1, -1, -1] + + +@pytest.mark.parametrize("skip_pool_shared", [False, True], ids=["fixed", "kill_switch"]) +def test_two_graph_keys_over_one_block_keep_their_uploads(cache, monkeypatch, skip_pool_shared): + monkeypatch.setattr(h2d_cache, "SKIP_POOL_SHARED_UPLOADS", skip_pool_shared) + key_a, key_b = _graph_key_views("msa_block_table") + replay_a = _capture_live_row_write(key_a, 4, [7, 7, 7, 7]) + table = [1, 2, 3, 4] + + assert copy_h2d_pool_shared(key_b[:4], _host(table)) is True + replay_a.replay() + copied = copy_h2d_pool_shared(key_b[:4], _host(table)) + torch.cuda.synchronize() + if skip_pool_shared: + assert copied is False + assert key_b[:4].tolist() == [7, 7, 7, 7] + else: + assert copied is True + assert key_b[:4].tolist() == table + + +def test_pool_shared_writes_drop_the_shadow_and_tail_record(cache): + buf = _device_zeros(6) + values = _host([1, 2]) + assert copy_h2d_if_changed(buf[:2], values) is True + fill_tail_if_changed(buf, 2, -1) + copy_h2d_pool_shared(buf[:2], _host([5, 6])) + buf[3] = 9 + # Neither record vouches for the buffer after a write outside the skip cache. + assert copy_h2d_if_changed(buf[:2], values) is True + fill_tail_if_changed(buf, 2, -1) + torch.cuda.synchronize() + assert buf.tolist() == [1, 2, -1, -1, -1, -1] + + +def test_pool_shared_writes_without_an_active_cache_are_plain_writes(): + assert not h2d_cache.enabled() + buf = _device_zeros(4) + assert copy_h2d_pool_shared(buf[:2], _host([3, 4]), pack=True) is True + fill_tail_pool_shared(buf, 2, -1) + torch.cuda.synchronize() + assert buf.tolist() == [3, 4, -1, -1] diff --git a/tests/unittest/_torch/executor/test_h2d_changed_span.py b/tests/unittest/_torch/executor/test_h2d_changed_span.py new file mode 100644 index 000000000000..2a987a260072 --- /dev/null +++ b/tests/unittest/_torch/executor/test_h2d_changed_span.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""H2DSkipCache(changed_span=True): a changed packed 2-D upload stages only each row's changed span.""" + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache, copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +ROWS, WIDTH = 4, 1024 + + +@pytest.fixture +def pack(): + return StepInputPack(torch.device("cuda"), parallel_scatter=True) + + +def _table(seed, rows=ROWS, width=WIDTH): + generator = torch.Generator().manual_seed(seed) + return torch.randint( + 0, 1 << 20, (rows, width), generator=generator, dtype=torch.int32 + ).pin_memory() + + +def _upload(pack, cache, dst, src): + """One step: open the pack, route the copy through the cache, issue the pack; returns the staged copies.""" + with h2d_cache.active(cache), step_inputs.active(pack): + step_inputs.begin() + copied = copy_h2d_if_changed(dst, src, pack=True) + staged = list(pack.h2d) + step_inputs.flush() + torch.cuda.synchronize() + return copied, staged + + +@pytest.mark.parametrize("changed_span", [False, True]) +def test_device_rows_equal_the_source_after_every_step(pack, changed_span): + """Bitwise: whatever is staged, the device view holds the source after each step.""" + cache = H2DSkipCache(changed_span=changed_span) + buf = torch.full((8, 2 * WIDTH), -7, dtype=torch.int32, device="cuda") + dst = buf[:ROWS, :WIDTH] + src = _table(0) + generator = torch.Generator().manual_seed(1) + for step in range(40): + src = src.clone().pin_memory() + # A request crosses into a new page: one entry of one row changes (sometimes several, in several rows). + for _ in range(1 + step % 3): + row = int(torch.randint(0, ROWS, (1,), generator=generator)) + col = int(torch.randint(0, WIDTH, (1,), generator=generator)) + src[row, col] = step + _upload(pack, cache, dst, src) + assert torch.equal(dst.cpu(), src) + # Rows and columns outside the view were never written. + assert bool((buf[ROWS:] == -7).all()) and bool((buf[:, WIDTH:] == -7).all()) + + +def test_one_changed_entry_stages_one_element(pack): + cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(2) + _upload(pack, cache, dst, src) + + src = src.clone().pin_memory() + src[2, 517] += 1 + copied, staged = _upload(pack, cache, dst, src) + + assert copied is True + assert staged == [(dst[2, 517].data_ptr(), staged[0][1], 4)] + assert torch.equal(dst.cpu(), src) + + +def test_each_changed_row_stages_its_first_to_last_changed_column(pack): + cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(3) + _upload(pack, cache, dst, src) + + src = src.clone().pin_memory() + src[0, 10] += 1 + src[0, 20] += 1 + src[3, 1000] += 1 + _, staged = _upload(pack, cache, dst, src) + + assert [(ptr, nbytes) for ptr, _, nbytes in staged] == [ + (dst[0, 10].data_ptr(), 11 * 4), + (dst[3, 1000].data_ptr(), 4), + ] + assert torch.equal(dst.cpu(), src) + + +def test_wider_view_uploads_the_columns_the_shadow_does_not_cover(pack): + cache = H2DSkipCache(changed_span=True) + buf = torch.full((ROWS, WIDTH), -1, dtype=torch.int32, device="cuda") + src = _table(4) + _upload(pack, cache, buf[:, : WIDTH // 2], src[:, : WIDTH // 2].contiguous().pin_memory()) + + _, staged = _upload( + pack, cache, buf[:, : WIDTH // 2 + 3], src[:, : WIDTH // 2 + 3].contiguous().pin_memory() + ) + + assert [nbytes for _, _, nbytes in staged] == [3 * 4] * ROWS + assert torch.equal(buf[:, : WIDTH // 2 + 3].cpu(), src[:, : WIDTH // 2 + 3]) + + +def test_a_mostly_changed_table_is_staged_whole(pack): + cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + _upload(pack, cache, dst, _table(5)) + + src = _table(6) + _, staged = _upload(pack, cache, dst, src) + + assert [nbytes for _, _, nbytes in staged] == [ROWS * WIDTH * 4] + assert torch.equal(dst.cpu(), src) + + +def test_off_by_default_the_whole_view_is_staged(pack): + cache = H2DSkipCache() + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(7) + _upload(pack, cache, dst, src) + + src = src.clone().pin_memory() + src[1, 1] += 1 + _, staged = _upload(pack, cache, dst, src) + + assert [nbytes for _, _, nbytes in staged] == [ROWS * WIDTH * 4] + + +def test_an_unpacked_copy_copies_the_whole_view(pack): + cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(8) + with h2d_cache.active(cache): + copy_h2d_if_changed(dst, src) + src = src.clone().pin_memory() + src[1, 1] += 1 + assert copy_h2d_if_changed(dst, src) is True + torch.cuda.synchronize() + assert torch.equal(dst.cpu(), src) + + +def test_a_view_with_another_row_stride_is_staged_whole(pack): + """The shadow describes other device bytes: nothing of it may be trusted.""" + cache = H2DSkipCache(changed_span=True) + buf = torch.zeros((2 * ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(9) + _upload(pack, cache, buf[:ROWS], src) + + _, staged = _upload(pack, cache, buf[::2], src) + + assert [nbytes for _, _, nbytes in staged] == [WIDTH * 4] * ROWS + assert torch.equal(buf[::2].cpu(), src) diff --git a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py index a8263f4df93e..baeddcc58479 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py +++ b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py @@ -491,6 +491,7 @@ def _preemption_executor() -> PyExecutor: # Wide enough that a mid-stream request emits nothing of its own, so any # response seen below came from the finish path rather than from streaming. executor.stream_interval = 100 + executor.first_content_emit_iters = 1 return executor diff --git a/tests/unittest/_torch/executor/test_kv_lens_update_once.py b/tests/unittest/_torch/executor/test_kv_lens_update_once.py new file mode 100644 index 000000000000..bee04e9db6df --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_lens_update_once.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RuntimeHooks.kv_lens_update_once: DecoderRunner._preprocess_inputs re-derives the KV-length metadata once, after the +overlap scheduler's KV-length correction, on exactly the steps that correct it, and leaves the inputs as before.""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks + + +class _Metadata: + """The fields _preprocess_inputs reads; on_update_kv_lens records the kv_lens it saw.""" + + def __init__(self, gen: int, tokens_per_gen: int, with_kv_cache: bool = True): + self.kv_cache_manager = object() if with_kv_cache else None + self.num_seqs = gen + self.num_contexts = 0 + self.num_generations = gen + self.num_ctx_tokens = 0 + self.num_chunked_ctx_requests = 0 + self.kv_lens_cuda = torch.arange(100, 100 + gen, dtype=torch.int32) + self.seen: list[list[int]] = [] + + def on_update_kv_lens(self) -> None: + self.seen.append(self.kv_lens_cuda.tolist()) + + +def _runner(once: bool, disable_overlap: bool = False) -> DecoderRunner: + r = object.__new__(DecoderRunner) + r.runtime_hooks = RuntimeHooks(kv_lens_update_once=once) + r._config = types.SimpleNamespace(disable_overlap_scheduler=disable_overlap) + r.previous_pos_id_offsets_cuda = torch.full((64,), -2, dtype=torch.int32) + r.previous_kv_lens_offsets_cuda = torch.full((16,), -3, dtype=torch.int32) + r.mapping = types.SimpleNamespace(has_cp_helix=lambda: False) + r.guided_decoder = None + return r + + +def _inputs(md: _Metadata, tokens: int) -> dict: + return { + "attn_metadata": md, + "input_ids": torch.zeros(tokens, dtype=torch.int32), + "position_ids": torch.arange(tokens, dtype=torch.int32).view(1, -1), + } + + +@pytest.mark.parametrize("gen", [1, 4]) +def test_corrected_step_updates_once_with_the_same_inputs(gen) -> None: + t = 4 + outs = {} + for once in (False, True): + md = _Metadata(gen, t) + inputs = _runner(once)._preprocess_inputs( + _inputs(md, gen * t), enable_spec_decode=True, runtime_draft_len=t - 1 + ) + outs[once] = (md.seen, inputs["position_ids"].tolist(), md.kv_lens_cuda.tolist()) + corrected = [100 + i - 3 for i in range(gen)] + # Without the hook: once before the correction, once after it. + assert outs[False][0] == [[100 + i for i in range(gen)], corrected] + # With it: only the call after the correction, which sees the same lengths. + assert outs[True][0] == [corrected] + assert outs[True][1:] == outs[False][1:] + + +@pytest.mark.parametrize( + "spec, disable_overlap, with_kv_cache", + [(False, False, True), (True, True, True), (True, False, False)], + ids=["no spec decode", "no overlap scheduler", "no kv cache manager"], +) +def test_uncorrected_steps_keep_their_update(spec, disable_overlap, with_kv_cache) -> None: + """Steps that do not correct the lengths still re-derive them once (the per-forward invalidation), with the same + inputs as without the hook.""" + outs = {} + for once in (False, True): + md = _Metadata(2, 4, with_kv_cache) + inputs = _runner(once, disable_overlap)._preprocess_inputs( + _inputs(md, 8), enable_spec_decode=spec, runtime_draft_len=3 + ) + outs[once] = (md.seen, inputs["position_ids"].tolist(), md.kv_lens_cuda.tolist()) + assert len(outs[True][0]) == 1 + assert outs[True] == outs[False] + + +def test_off_by_default() -> None: + assert RuntimeHooks().kv_lens_update_once is False diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py index 2dadaa957e8e..37a175ffffb7 100644 --- a/tests/unittest/_torch/executor/test_runtime_hooks.py +++ b/tests/unittest/_torch/executor/test_runtime_hooks.py @@ -34,8 +34,17 @@ def test_model_without_hooks_gets_the_all_off_default(): assert not hooks.skip_unchanged_h2d assert not hooks.loop_on_execution_stream assert not hooks.merge_tuned_tactics_across_ranks + assert hooks.first_content_emit_iters == 1 + assert not hooks.mxfp8_pinned_tactics + assert not hooks.late_compile_detect assert not hooks.skip_greedy_rng_upload + assert not hooks.pack_spec_rng_uploads assert hooks.spec_step_fold is None + assert not hooks.sampler_d2h_side_stream + assert not hooks.spec_small_h2d + assert not hooks.iteration_timing_one_event + assert not hooks.pack_parallel_scatter + assert not hooks.pack_in_graph assert hooks.msa == MsaRuntimePolicy() assert not hooks.msa.short_sequence_identity assert hooks.msa.kv_update is None diff --git a/tests/unittest/_torch/executor/test_runtime_opt_ins.py b/tests/unittest/_torch/executor/test_runtime_opt_ins.py new file mode 100644 index 000000000000..4f3d8b2eec74 --- /dev/null +++ b/tests/unittest/_torch/executor/test_runtime_opt_ins.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The process-wide features a model opts into through RuntimeHooks and the executor arms +(``_arm_runtime_opt_ins``): the SM107 MXFP8 GEMM tactic table and the late-compile detector. A model +that declares neither leaves both off, even at a table shape on SM107; the MiniMax-M3 sm_107 targets +declare both.""" + +import importlib + +import pytest +import torch + +from tensorrt_llm._torch import late_compile +from tensorrt_llm._torch.custom_ops import mxfp8_cute_pinned_tactics as pin +from tensorrt_llm._torch.pyexecutor.py_executor import _arm_runtime_opt_ins +from tensorrt_llm._torch.pyexecutor.runtime_hooks import runtime_hooks_of + +pytestmark = pytest.mark.cpu_only + +M3_SHAPE = (64, 6144, 6144) # (m, n, k): a table entry +M3_TARGETS = "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__{}.modeling" + + +class _PlainModel(torch.nn.Module): + """A model that declares no hooks.""" + + +@pytest.fixture(autouse=True) +def _sm107_unarmed(monkeypatch): + monkeypatch.setattr(pin, "_ARMED", [False]) + monkeypatch.setattr(pin, "ENABLED", True) + monkeypatch.setattr(pin, "get_sm_version", lambda: 107) + monkeypatch.setattr(late_compile, "_ready", [False]) + + +def test_a_model_without_the_opt_in_arms_nothing_at_an_m3_shape(): + hooks = runtime_hooks_of(_PlainModel()) + + assert _arm_runtime_opt_ins(hooks) is False + assert pin.tactic_for(*M3_SHAPE) is None + assert not late_compile._ready[0] + + +@pytest.mark.parametrize("tp", ["tp2", "tp4", "tp8"]) +def test_the_m3_sm107_targets_arm_both(tp): + hooks = importlib.import_module(M3_TARGETS.format(tp))._runtime_hooks(None) + + assert hooks.mxfp8_pinned_tactics and hooks.late_compile_detect + assert _arm_runtime_opt_ins(hooks) is True + m, n, k = M3_SHAPE + assert pin.tactic_for(m, n, k) == pin.TACTICS[(n, k, pin._bucket(m))] + + +def test_the_kill_switch_still_wins_once_armed(monkeypatch): + pin.arm() + monkeypatch.setattr(pin, "ENABLED", False) + + assert pin.tactic_for(*M3_SHAPE) is None diff --git a/tests/unittest/_torch/executor/test_step_inputs.py b/tests/unittest/_torch/executor/test_step_inputs.py index c1b4b95a02d4..9ca87ba292fd 100644 --- a/tests/unittest/_torch/executor/test_step_inputs.py +++ b/tests/unittest/_torch/executor/test_step_inputs.py @@ -602,3 +602,343 @@ def worker(index): assert results == {0: (True, 1), 1: (True, 1)} assert dsts[0].tolist() == [1] * 4 assert dsts[1].tolist() == [2] * 4 + + +def _random_pack_run(parallel: bool, seed: int, overlap: bool): + """Stage a fixed random mix of host copies (2-D page-table rows included) and device copies that read a host + copy's destination; return every destination after the flush.""" + g = torch.Generator().manual_seed(seed) + pack = StepInputPack(_device(), parallel_scatter=parallel) + big = _zeros(4096) + table = torch.zeros((8, 64), dtype=torch.int32, device=_device()) + out = _zeros(256) + with step_inputs.active(pack): + step_inputs.begin() + for k in range(24): + n = int(torch.randint(1, 64, (1,), generator=g)) + step_inputs.copy_( + big[k * 128 : k * 128 + n], + torch.randint(0, 1 << 30, (n,), generator=g, dtype=torch.int32), + ) + step_inputs.copy_( + table[:3, :40], torch.randint(0, 1 << 30, (3, 40), generator=g, dtype=torch.int32) + ) + if overlap: + # A later host copy overwrites part of an earlier one: order matters, so the serial scatter must run. + step_inputs.copy_(big[0:4], torch.full((4,), -7, dtype=torch.int32)) + step_inputs.copy_(out[:64], big[:64]) + step_inputs.copy_(out[64:96], big[128:160]) + step_inputs.flush() + torch.cuda.synchronize() + return big.cpu(), table.cpu(), out.cpu() + + +@pytest.mark.parametrize("overlap", [False, True], ids=["disjoint", "overlapping"]) +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_parallel_scatter_writes_what_the_serial_scatter_writes(seed, overlap): + serial = _random_pack_run(False, seed, overlap) + parallel = _random_pack_run(True, seed, overlap) + for a, b in zip(serial, parallel): + assert torch.equal(a, b) + + +def _compiled_variants(kernel) -> int: + """Compiled variants of a Triton kernel on this device, or -1 where this Triton keeps them elsewhere.""" + caches = getattr(kernel, "device_caches", None) + if not caches: + return -1 + cache = caches.get(torch.cuda.current_device()) + return len(cache[0]) if cache else 0 + + +@pytest.mark.parametrize("count", [1, 15, 16, 17, 120]) # 120: 240 of the 256 descriptors +def test_parallel_scatter_is_bitwise_for_every_copy_count(count): + """``first`` (the device copies' descriptor offset) is the step's host-copy count: one compiled variant serves + every count, so a new count never compiles mid-serving.""" + pack = StepInputPack(_device(), parallel_scatter=True) + width = 3 + dst = _zeros(count * width) + mirror = _zeros(count * width) + g = torch.Generator().manual_seed(count) + values = torch.randint( + -(1 << 31), (1 << 31) - 1, (count, width), generator=g, dtype=torch.int32 + ) + with step_inputs.active(pack): + step_inputs.begin() + for i in range(count): + step_inputs.copy_(dst[i * width : (i + 1) * width], values[i]) + for i in range(count): + step_inputs.copy_(mirror[i * width : (i + 1) * width], dst[i * width : (i + 1) * width]) + assert len(pack.h2d) == count and len(pack.d2d) == count + step_inputs.flush() + torch.cuda.synchronize() + assert torch.equal(dst.cpu(), values.flatten()) + assert torch.equal(mirror.cpu(), values.flatten()) + variants = _compiled_variants(step_inputs._parallel_scatter_kernel) + assert variants in (-1, 1, 2), variants # at most one per DEVICE_SRC + + +def test_parallel_scatter_only_for_order_free_packs(): + assert step_inputs._parallel_ok([(0, 0, 8), (8, 0, 8)], []) + assert not step_inputs._parallel_ok([(0, 0, 8), (4, 0, 8)], []) + # A device copy may read a host copy's destination (the launches run in order) ... + assert step_inputs._parallel_ok([(0, 0, 8)], [(100, 0, 8)]) + # ... but not another device copy's destination, and device destinations must be disjoint. + assert not step_inputs._parallel_ok([], [(100, 0, 8), (200, 100, 8)]) + assert not step_inputs._parallel_ok([], [(100, 0, 8), (104, 300, 8)]) + + +def test_parallel_scatter_is_off_by_default(): + assert not StepInputPack(_device()).parallel_scatter + + +def _graph_case(early=False, epoch=False): + """A pack with the in-graph prologue and a graph that starts with it and then reads the destinations.""" + pack = StepInputPack( + _device(), parallel_scatter=True, in_graph=True, early_copy=early, step_epoch=epoch + ) + a, b, c = _zeros(16), _zeros(16), _zeros(16) + out = _zeros(48) + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, _host(values[0])) + step_inputs.copy_(b, _host(values[1])) + if overlap: + step_inputs.copy_(a[:4], _host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if flush: + step_inputs.flush(in_graph=True) + else: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_in_graph_pack_is_issued_by_the_replay(): + pack, graph, dsts, out = _graph_case() + for step in range(5): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + _step(pack, graph, dsts, (va, vb)) + torch.cuda.synchronize() + assert out.tolist() == va + vb + va, step + # Nothing was issued outside the graph: the window carried the copies. + assert int(pack.ghost[:8].view(torch.int64)[0]) == 2 + + +def test_in_graph_pack_not_order_free_is_issued_eagerly(): + pack, graph, dsts, out = _graph_case() + va = list(range(16)) + vb = list(range(100, 116)) + _step(pack, graph, dsts, (va, vb), overlap=True) + torch.cuda.synchronize() + expect_a = [-1, -2, -3, -4] + va[4:] + assert out.tolist() == expect_a + vb + expect_a + # The window was zeroed, so the prologue wrote nothing over the eager result. + assert not pack.ghost[:16].any() + + +def test_in_graph_replay_without_a_pack_reapplies_nothing(): + pack, graph, dsts, out = _graph_case() + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + _step(pack, graph, dsts, ([0] * 16, [0] * 16), flush=False) + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_pack_lands_every_step(early): + pack, graph, dsts, out = _graph_case(early) + assert pack.early_copy is early + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + _step(pack, graph, dsts, (va, vb), overlap=step % 2 == 1) + torch.cuda.synchronize() + head = [-1, -2, -3, -4] if step % 2 == 1 else [] + expect_a = head + va[len(head) :] + assert out.tolist() == expect_a + vb + expect_a, step + + +def _history_run(pack, graph, dsts, out, steps, before_replay=None, before_upload=None): + """Steps back to back without host syncs; each replay's output is copied to a history row behind it (same + stream), so a later upload that overwrote the window too early shows up as the wrong step's values.""" + hist = _zeros(steps, out.numel()) + expect = [] + for step in range(steps): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + expect.append(va + vb + va) + a, b, c = dsts + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, _host(va)) + step_inputs.copy_(b, _host(vb)) + step_inputs.copy_(c, a) + if before_upload is not None: + before_upload(step) + step_inputs.flush(in_graph=True) + if before_replay is not None: + before_replay(step) + graph.replay() + hist[step].copy_(out) + torch.cuda.synchronize() + return hist.tolist(), expect + + +def test_early_copy_waits_for_the_previous_prologue_before_overwriting_the_window(): + """A slow GPU: the replay of step N (and its prologue) is held back behind a long kernel while the host already + uploads step N+1; that upload must wait for step N's prologue to consume the window.""" + pack, graph, dsts, out = _graph_case(early=True) + hist, expect = _history_run( + pack, graph, dsts, out, 4, before_replay=lambda step: torch.cuda._sleep(_SLEEP_CYCLES // 4) + ) + assert hist == expect + + +def test_early_copy_prologue_waits_for_a_slow_upload(): + """A slow copy stream: each step's upload is held back behind a long kernel on the copy stream; the prologue must + wait for it (staged_ready), and the next host fill must wait for the upload to have read the window.""" + pack, graph, dsts, out = _graph_case(early=True) + + def slow_copy(step): + with torch.cuda.stream(pack.copy_stream): + torch.cuda._sleep(_SLEEP_CYCLES // 4) + + hist, expect = _history_run(pack, graph, dsts, out, 4, before_upload=slow_copy) + assert hist == expect + + +def test_early_copy_with_a_slow_host(): + """A slow host: the GPU idles between steps (the window lands long before the replay).""" + import time + + pack, graph, dsts, out = _graph_case(early=True) + hist, expect = _history_run( + pack, graph, dsts, out, 3, before_replay=lambda step: time.sleep(0.05) + ) + assert hist == expect + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_replay_after_quiesce_reapplies_nothing(early): + pack, graph, dsts, out = _graph_case(early) + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +def test_early_copy_needs_in_graph(): + assert not StepInputPack(_device(), parallel_scatter=True, early_copy=True).early_copy + + +def test_prologue_op_needs_early_copy_and_is_off_by_default(): + def op(window, epoch, max_desc): + raise AssertionError("not called") + + assert StepInputPack(_device(), parallel_scatter=True, early_copy=True).prologue_op is None + assert ( + StepInputPack(_device(), parallel_scatter=True, in_graph=True, prologue_op=op).prologue_op + is None + ) + + +def test_in_graph_needs_parallel_scatter_and_is_off_by_default(): + assert not StepInputPack(_device()).in_graph + assert not StepInputPack(_device(), in_graph=True).in_graph + + +def test_in_graph_replay_outside_the_pack_context_reapplies_nothing(): + pack, graph, dsts, out = _graph_case() + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + # What the graph runner does for a replay outside the pack's context. + assert not step_inputs.is_active(pack) + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +def test_step_epoch_is_off_by_default_and_needs_in_graph(): + assert StepInputPack(_device(), parallel_scatter=True, in_graph=True).epoch is None + assert StepInputPack(_device(), parallel_scatter=True, step_epoch=True).epoch is None + with step_inputs.active(StepInputPack(_device(), parallel_scatter=True, in_graph=True)): + assert step_inputs.step_epoch() is None + assert step_inputs.step_epoch() is None + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_step_epoch_advances_once_per_replay_and_never_resets(early): + pack, graph, dsts, out = _graph_case(early, epoch=True) + with step_inputs.active(pack): + assert step_inputs.step_epoch() is pack.epoch + torch.cuda.synchronize() + start = int(pack.epoch.item()) # the warm-up compile launched the prologue once + assert start == 1 + for step in range(5): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + # A replay with a new pack, one with an eagerly issued pack, and one with nothing to scatter. + _step(pack, graph, dsts, (va, vb), overlap=step == 2, flush=step != 3) + torch.cuda.synchronize() + assert int(pack.epoch.item()) == start + step + 1, step + # The copies are unchanged by the epoch. + _step(pack, graph, dsts, (list(range(16)), list(range(16, 32)))) + torch.cuda.synchronize() + assert out.tolist() == list(range(16)) + list(range(16, 32)) + list(range(16)) + assert int(pack.epoch.item()) == start + 6 + + +def test_step_epoch_is_ordered_before_the_graph_body(): + """A kernel after the prologue in the same graph reads the bumped epoch of its own replay.""" + pack = StepInputPack(_device(), parallel_scatter=True, in_graph=True, step_epoch=True) + seen = torch.zeros(4, dtype=torch.int64, device="cuda") + slot = torch.zeros(1, dtype=torch.int64, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + step_inputs.capture_prologue() + seen.index_copy_(0, slot, pack.epoch) + slot.add_(1) + torch.cuda.synchronize() + start = int(pack.epoch.item()) + for _ in range(4): + with step_inputs.active(pack): + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert seen.tolist() == [start + 1, start + 2, start + 3, start + 4] + + +def test_eager_flush_does_not_advance_the_step_epoch(): + pack = StepInputPack(_device(), parallel_scatter=True, in_graph=True, step_epoch=True) + torch.cuda.synchronize() + start = int(pack.epoch.item()) + dst = _zeros(4) + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(dst, _host([1, 2, 3, 4])) + step_inputs.flush() + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 4] + assert int(pack.epoch.item()) == start diff --git a/tests/unittest/_torch/kernels/parallel/test_mxfp8_mxfp8_gemm.py b/tests/unittest/_torch/kernels/parallel/test_mxfp8_mxfp8_gemm.py index d6b489dad7b1..13b2e935d4c0 100644 --- a/tests/unittest/_torch/kernels/parallel/test_mxfp8_mxfp8_gemm.py +++ b/tests/unittest/_torch/kernels/parallel/test_mxfp8_mxfp8_gemm.py @@ -23,8 +23,9 @@ @pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="MXFP8 GEMM is supported on SM100 and SM103 only. Current SM is %d." % getSMVersion(), + getSMVersion() not in (100, 103, 107), + reason="MXFP8 GEMM is supported on SM100, SM103 and SM107 only. Current SM is %d." + % getSMVersion(), ) @pytest.mark.parametrize( "m,n,k", @@ -62,8 +63,8 @@ def test_mxfp8_mxfp8_gemm_large_m(m: int, n: int, k: int): @pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="MXFP8 tactic runner requires SM100 or SM103. Current SM is %d." % getSMVersion(), + getSMVersion() not in (100, 103, 107), + reason="MXFP8 tactic runner requires SM100, SM103 or SM107. Current SM is %d." % getSMVersion(), ) def test_mxfp8_mxfp8_runner_tactics(): """Every exposed tactic and the generic fallback produce aligned output.""" @@ -94,8 +95,8 @@ def test_mxfp8_mxfp8_runner_tactics(): @pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="MXFP8 tactic cache requires SM100 or SM103. Current SM is %d." % getSMVersion(), + getSMVersion() not in (100, 103, 107), + reason="MXFP8 tactic cache requires SM100, SM103 or SM107. Current SM is %d." % getSMVersion(), ) def test_mxfp8_mxfp8_native_tactic_cache(): """The direct op uses cached tactics and preserves the generic fallback.""" @@ -156,8 +157,8 @@ def test_mxfp8_mxfp8_native_tactic_cache(): @pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="MXFP8 tactic cache requires SM100 or SM103. Current SM is %d." % getSMVersion(), + getSMVersion() not in (100, 103, 107), + reason="MXFP8 tactic cache requires SM100, SM103 or SM107. Current SM is %d." % getSMVersion(), ) @pytest.mark.parametrize( "lower_bound,upper_bound,bucket", diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..4052191ded74 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_silu_and_mul entry of the MiniMax-M3 tp2 target: SwiGLU-OAI and plain SwiGLU.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.activation.m3_silu_and_mul import ( # noqa: E501 + m3_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + # Gate and up spread past the clamp (|x| up to ~12) so both clamps engage. + x = (torch.randn(spec["m"], 2 * spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + return (x, *OAI) if spec["oai"] else (x, None, None, None) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few + fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); + measured on GB300: worst element 0.5 unit, relative RMS 0.01.""" + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(*args) + assert out.shape == (cell.spec["m"], cell.spec["inter"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=1.0, rms_ulp=0.1) + + +def test_the_gate_sees_a_missing_clamp() -> None: + """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" + spec = op.CELLS[0].spec + x, *oai = _build(dict(spec, m=64), 3) + out = op(x, *oai) + with pytest.raises(AssertionError): + assert_within_ulp(out, op.reference(x, None, oai[1], oai[2]), element_ulp=1.0, rms_ulp=0.1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_attn_runtime_tp2.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_attn_runtime_tp2.py new file mode 100644 index 000000000000..137230cd0345 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_attn_runtime_tp2.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real MiniMax-M3 attention runtime state at the TP2 per-rank geometry, for the backend-step entries' GPU tests. + +One hybrid NVFP4 KV cache manager (``MiniMaxM3KVCacheManagerV2``) laid out as the target's: a dense layer in the FP8 +half, a sparse MSA layer in the NVFP4 half with its FP8 index-K cache, and (``draft=True``) the one-model Eagle3 +drafter's shared layer, read through the manager's FP8 draft view as the runtime's draft swap reads it. Metadata is +the runtime's own type, filled by its own ``prepare()``; the backends are built the way the target builds them. + +The geometry is one rank's: 32 query heads over two KV heads and two index heads of dim 128 (one index-K head). +The mapping is a single +rank holding exactly those heads, so nothing is sharded again. + +History rows are written by the in-tree scatters (``fused_write_layer_caches`` for the FP8 layers, +``fused_write_layer_caches_nvfp4`` for the NVFP4 layer and its index-K), at the slots the manager assigned, which are +not in request order. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + MiniMaxM3MsaSparseAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, + fused_write_layer_caches_nvfp4, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HEADS = 32 +KV_HEADS = 2 +INDEX_HEADS = 2 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +DENSE, SPARSE, DRAFT = 0, 1, 2 +MAX_BATCH = 8 +MAX_SEQ_LEN = 8192 +MAX_NUM_TOKENS = 8192 +#: The Eagle3 drafter checkpoint's RoPE: NeoX over the whole head, theta 5e6. +DRAFT_ROPE_THETA = 5.0e6 + +#: E4M3-exact planted index levels, 1/8 apart in relative terms (0.25 .. 15): 48 distinct block scores. +LEVELS = tuple(2.0**e * (1 + m / 8) for e in range(-2, 4) for m in range(8)) +#: The index value of a block's non-peak rows; its score (8) is below every level's (>= 16). +LOW = 0.125 + + +def mapping() -> Mapping: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + +def sparse_attention_config() -> MiniMaxM3SparseAttentionConfig: + """The target's runtime: MSA, the FP8 indexer, the fused five-way projection (M3's checkpoint selection rule: + no forced initial blocks, the query's own block forced).""" + return MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_num_index_heads=INDEX_HEADS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + ) + + +def make_manager(*, draft: bool) -> MiniMaxM3KVCacheManagerV2: + torch.cuda.init() + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_tokens=4 * MAX_SEQ_LEN, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + event_buffer_max_size=0, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=[KV_HEADS] * (3 if draft else 2), + head_dim=HEAD_DIM, + tokens_per_block=PAGE, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + mapping=mapping(), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") if draft else None + ), + sparse_layer_ids=[SPARSE], + sparse_attention_config=sparse_attention_config(), + ) + + +def make_msa_backend(layer: int) -> MiniMaxM3MsaSparseAttention: + """One layer's MSA backend as ``MiniMaxM3Tp2ModelingV2`` builds it (NVFP4 KV quant config on every layer).""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + backend = create_attention( + "TRTLLM", + layer, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=None, + quant_config=quant, + q_scaling=1.0, + sparse_params=sparse_attention_config().to_sparse_params(layer_idx=layer), + ) + backend.update_quant_config(quant) + assert isinstance(backend, MiniMaxM3MsaSparseAttention), type(backend).__name__ + return backend + + +def make_draft_backend() -> TrtllmAttention: + """The drafter layer's attention as ``M3Eagle3Draft`` builds it: fused NeoX RoPE, an FP8 KV cache (the target + overrides the drafter's KV dtype to FP8), no other quantization.""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + backend = create_attention( + "TRTLLM", + DRAFT, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=HEAD_DIM, theta=DRAFT_ROPE_THETA, max_positions=MAX_SEQ_LEN), + ), + quant_config=quant, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + backend.update_quant_config(quant) + assert type(backend) is TrtllmAttention, type(backend).__name__ + return backend + + +def make_msa_metadata(manager) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = MiniMaxM3MsaSparseAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + max_num_sequences=MAX_BATCH, + kv_cache_manager=manager, + mapping=mapping(), + sparse_metadata_params=sparse_attention_config().to_sparse_metadata_params(), + ) + # The target's MSA preparation choice that changes what prepare() stages (the identity tables of short steps). + metadata.msa_policy = MsaRuntimePolicy(short_sequence_identity=True) + return metadata + + +def make_draft_metadata(manager) -> TrtllmAttentionMetadata: + """Drafter metadata over the shared manager's FP8 draft view, the manager the runtime's draft swap installs.""" + view = manager.get_draft_kv_cache_view() + assert view is not None and view.dtype == DataType.FP8 + return TrtllmAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=mapping(), + ) + + +def allocate(manager, request_ids: Sequence[int], qlens: Sequence[int], cached: Sequence[int]): + manager.add_dummy_requests(list(request_ids), token_nums=[c + n for c, n in zip(cached, qlens)]) + + +def prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + c + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +def flat_slots(manager, layer: int, request_id: int, positions: torch.Tensor) -> torch.Tensor: + """The cache slots (page slot id * PAGE + offset) of a request's positions in ``layer``.""" + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0], + dtype=torch.int64, + device="cuda", + ) + positions = positions.to(device="cuda", dtype=torch.int64) + return (pages[positions // PAGE] * PAGE + positions % PAGE).to(torch.int32) + + +def write_fp8_rows(manager, layer, request_id, positions, k, v) -> None: + """K / V rows [n, 128] (E4M3) into an FP8 layer's cache at ``positions``.""" + view = manager.get_buffers(layer, "HND") + assert view.dtype == torch.float8_e4m3fn + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches(view[:, 0], view[:, 1], None, slots, k, v, None) + + +def write_nvfp4_rows(manager, metadata, layer, request_id, positions, k, v, idx_k, quant) -> None: + """BF16 K / V rows quantized into the NVFP4 layer's cache, and index-K rows into its FP8 index cache.""" + data = manager.get_buffers(layer, "HND") + scales = manager.get_block_scale_buffers(layer, "HND") + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches_nvfp4( + data[:, 0], + data[:, 1], + scales[:, 0], + scales[:, 1], + metadata.msa_idx_k_cache(layer), + slots, + k, + v, + idx_k, + quant, + ) + + +def index_directions(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """Two orthogonal ±1 index directions (the halves of the 128 dims), E4M3-exact.""" + signs = torch.randint(0, 2, (HEAD_DIM,), generator=gen, device="cuda").float() * 2 - 1 + u0 = torch.zeros(HEAD_DIM, device="cuda") + u1 = torch.zeros(HEAD_DIM, device="cuda") + u0[: HEAD_DIM // 2] = signs[: HEAD_DIM // 2] + u1[HEAD_DIM // 2 :] = signs[HEAD_DIM // 2 :] + return u0, u1 + + +def planted_index_keys(length: int, u0, u1, gen: torch.Generator) -> torch.Tensor: + """Index-K rows [length, 128] (bf16, E4M3-exact) whose per-block maximum score is unambiguous. + + Every block has one peak row at a random position, ``a_b u0 + c_b u1`` with ``a`` and ``c`` two independent + permutations of ``LEVELS``; its other rows are ``LOW (u0 + u1)``. A query along ``u0`` then ranks the blocks by + ``a``, one along ``u1`` by ``c``, every score 1/8 or more apart from the next.""" + blocks = (length + PAGE - 1) // PAGE + assert blocks <= len(LEVELS) + levels = torch.tensor(LEVELS, device="cuda") + a = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + c = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + rows = (LOW * (u0 + u1)).expand(length, HEAD_DIM).clone() + for b in range(blocks): + span = min(PAGE, length - b * PAGE) + peak = b * PAGE + int(torch.randint(0, span, (1,), generator=gen, device="cuda")) + rows[peak] = a[b] * u0 + c[b] * u1 + return rows.to(torch.bfloat16) + + +def nvfp4_quant(*values: torch.Tensor) -> torch.Tensor: + """The (Q, K, V) quantization scales ``fused_write_layer_caches_nvfp4`` takes: each role's amax at the top of the + E2M1 x E4M3 range.""" + quant = torch.ones(3, dtype=torch.float32, device="cuda") + for role, x in enumerate(values, start=1): + quant[role] = (448.0 * 6.0) / x.abs().max().float() + return quant diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_nvfp4_cache.py new file mode 100644 index 000000000000..61886497be86 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_nvfp4_cache.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A MiniMax-M3 TP2 sparse-layer decode step over a packed NVFP4 cache, for the attention entries' tests. + +The cache is written by the in-tree scatter (``fused_write_layer_caches_nvfp4``, +which quantizes with ``torch.ops.trtllm.fp4_quantize``), into one +``[pages, 2 (K, V), 2 (KV heads), 128, .]`` pool per role exactly as ``msa_paged_kv`` and +``get_block_scale_buffers`` hand the K and V views over. Pages are shuffled so +a kernel that ignored the block table would read the wrong ones. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +NUM_HEADS = 32 +NUM_KV_HEADS = 2 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 + + +def token_kv_len(seq_len: int, decode_query_len: int, offset: int) -> int: + """The KV extent query token ``offset`` of a request attends, as every decode kernel derives it.""" + return max(seq_len - decode_query_len + offset + 1, 0) + + +def build_step( + seq_lens: Sequence[int], decode_query_len: int, *, identity: bool, seed: int +) -> SimpleNamespace: + """One decode step: E4M3 q [T, 32, 128], the packed NVFP4 K/V views and their scales, a [2, T, 16] + head-major top-k table (the identity table, or a random ascending selection with a -1 tail per KV head), the + block table and per-request lengths.""" + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = max(1, max((s + PAGE - 1) // PAGE for s in seq_lens)) + num_pages = batch * max_blocks + num_slots = num_pages * PAGE + block_table = ( + torch.randperm(num_pages, device="cuda", generator=g) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + data_pool = torch.zeros( + num_pages, 2, NUM_KV_HEADS, PAGE, HEAD_DIM // 2, dtype=torch.uint8, device="cuda" + ) + scale_pool = torch.zeros( + num_pages, 2, NUM_KV_HEADS, PAGE, HEAD_DIM // 16, dtype=torch.uint8, device="cuda" + ) + quant = torch.ones(3, dtype=torch.float32, device="cuda") + values = [] + for role in (0, 1): + x = torch.randn( + num_slots, NUM_KV_HEADS * HEAD_DIM, device="cuda", generator=g, dtype=torch.float32 + ).to(torch.bfloat16) + quant[role + 1] = (448.0 * 6.0) / x.abs().max().float().item() + values.append(x) + assert fused_write_layer_caches_nvfp4( + data_pool[:, 0], + data_pool[:, 1], + scale_pool[:, 0], + scale_pool[:, 1], + None, + torch.arange(num_slots, device="cuda", dtype=torch.int32), + values[0], + values[1], + None, + quant, + ) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) + topk = torch.full((NUM_KV_HEADS, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + nblocks = (token_kv_len(seq_lens[req], decode_query_len, offset) + PAGE - 1) // PAGE + for head in range(NUM_KV_HEADS): + if identity: + assert nblocks <= TOPK, "an identity step has every request within 16 blocks" + chosen = torch.arange(nblocks, device="cuda") + else: + chosen = torch.randperm(nblocks, device="cuda", generator=g)[:TOPK].sort().values + topk[head, tok, : chosen.numel()] = chosen.to(torch.int32) + return SimpleNamespace( + q=q, + k_paged=data_pool[:, 0], + v_paged=data_pool[:, 1], + k_block_scale=scale_pool[:, 0], + v_block_scale=scale_pool[:, 1], + k_global_scale=(1.0 / quant[1:2]).contiguous(), + v_global_scale=(1.0 / quant[2:3]).contiguous(), + topk=topk, + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device="cuda"), + decode_query_len=decode_query_len, + sm_scale=HEAD_DIM**-0.5, + source=values, + quant=quant, + num_slots=num_slots, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py new file mode 100644 index 000000000000..1bd4e703e9dc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA page table is written on the host side only, before the graph replay. + +The index scorer's early issue (``kernels/m3_index_score.py``, ``early_issue``) reads ``block_table[b, split]`` BEFORE +its PDL grid-dependency wait and loads that page's K tile right after it. That is exact only while no kernel of the +step writes ``MsaMetadata.msa_block_table``: today its one writer is the host-to-device copy in +``_build_msa_fields`` (called from ``prepare()``, before the replay), besides the allocation. This test scans +``tensorrt_llm/_torch`` and fails on any other write site, so a new in-graph writer has to revisit the early issue. + +(``msa_seq_lens_cuda`` IS patched in-graph by ``on_update_kv_lens``; the scorer reads it before the wait only as a hint +and bounds every load and store by the value it reads after the wait.) + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +_NAME = "msa_block_table" +# In-place tensor methods that write their receiver. +_MUTATORS = { + "copy_", + "fill_", + "zero_", + "index_put_", + "index_copy_", + "scatter_", + "masked_fill_", + "add_", + "sub_", +} +# Helpers whose first argument is the destination they write. +_COPY_HELPERS = {"copy_h2d_if_changed", "copy_h2d_pool_shared"} +# (file relative to tensorrt_llm/_torch, enclosing function) -> what it may do +_ALLOWED = { + ( + "attention/backends/sparse/minimax_m3/msa_backend.py", + "_create_msa_buffers", + ), # the allocation + ( + "attention/backends/sparse/minimax_m3/msa_backend.py", + "_build_msa_fields", + ), # the H2D copy in prepare() +} + + +def _mentions(node) -> bool: + """Whether the expression's receiver chain (attributes / subscripts / slices) is ``*.msa_block_table``.""" + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute) and node.attr == _NAME: + return True + node = node.func if isinstance(node, ast.Call) else node.value + return isinstance(node, ast.Name) and node.id == _NAME + + +def _write_sites(tree: ast.AST): + sites = [] + + def visit(node, func): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + func = node.name + if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + if any(_mentions(t) for t in targets) and not isinstance(node, ast.AnnAssign): + sites.append((func, node.lineno)) + if isinstance(node, ast.Call): + f = node.func + if isinstance(f, ast.Attribute) and f.attr in _MUTATORS and _mentions(f.value): + sites.append((func, node.lineno)) + name = ( + f.id + if isinstance(f, ast.Name) + else (f.attr if isinstance(f, ast.Attribute) else "") + ) + if name in _COPY_HELPERS and node.args and _mentions(node.args[0]): + sites.append((func, node.lineno)) + for child in ast.iter_child_nodes(node): + visit(child, func) + + visit(tree, None) + return sites + + +def test_only_prepare_writes_the_msa_page_table() -> None: + found = set() + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if _NAME not in src: + continue + rel = path.relative_to(_ROOT).as_posix() + for func, line in _write_sites(ast.parse(src)): + assert (rel, func) in _ALLOWED, ( + f"{rel}:{line} ({func}) writes msa_block_table; the index scorer's early issue reads it before its " + "PDL grid wait (kernels/m3_index_score.py), so an in-graph writer must be ordered before it or the " + "early page read moved after the wait" + ) + found.add((rel, func)) + assert found == _ALLOWED, f"expected write sites {_ALLOWED}, found {found}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..0ebd8285aa9e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7): the TP2 dense layers' FP8 q/k/v (minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8) carries the bits +of the stock path, the BF16 fused_qk_norm_rope output cast to E4M3, for every norm / RoPE variant: Q and K are rounded +to BF16 before E4M3 (K is stored in the FP8 cache), V is cast.""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers the trtllm ops) + +assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) + +HQ, HK, HV, D, ROT = 32, 2, 2, 128, 64 + + +@pytest.mark.parametrize("tokens", [1, 4, 8, 4096]) +@pytest.mark.parametrize("use_gemma", [False, True], ids=["rmsnorm", "gemma"]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_dense_fp8_qkv_is_the_bf16_kernel_cast(tokens, use_gemma, is_neox) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8") + g = torch.Generator(device="cuda").manual_seed(tokens * 4 + int(use_gemma) * 2 + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + ours = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8( + qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, use_gemma + ) + stock = qkv.clone() + torch.ops.trtllm.fused_qk_norm_rope( + stock, HQ, HK, HV, D, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, + 1.0, 0.0, 0.0, 1.0, True, use_gemma, False, 0, 0, + ) # fmt: skip + assert torch.equal(ours.view(torch.uint8), stock.to(torch.float8_e4m3fn).view(torch.uint8)) + + +@pytest.mark.parametrize("tokens,live", [(1, 1), (4, 4), (8, 5), (64, 64)]) +@pytest.mark.parametrize("tokens_per_block", [32, 128]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_fused_kv_insert_stores_the_bytes_of_write_kv_slots_and_the_compact_q( + tokens, live, tokens_per_block, is_neox +) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + g = torch.Generator(device="cuda").manual_seed(tokens * 1000 + tokens_per_block + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + pages, layers, layer = 16, 3, 1 + # The layer's HND view of a multi-layer pool, strided as get_buffers returns it. + raw = torch.randint( + 0, + 256, + (pages, layers, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.randperm(pages * tokens_per_block, device="cuda", generator=g)[:tokens].to( + torch.int32 + ) + slots[live:] = -1 # padded rows own no slot + + def view(t): + return t.view(torch.float8_e4m3fn)[:, layer].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, False) + q_compact = torch.empty(tokens, HQ, D, device="cuda", dtype=torch.float8_e4m3fn) + ours = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, live, q_compact + ) + plain = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + assert torch.equal(ours.view(torch.uint8), plain.view(torch.uint8)) + assert torch.equal( + q_compact.view(tokens, HQ * D).view(torch.uint8), plain[:, : HQ * D].view(torch.uint8) + ) + + ref_raw = raw.clone() + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D) + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D) + ref = view(ref_raw) + write_kv_slots(ref[:, 0], slots, k, live, layout="HND") + write_kv_slots(ref[:, 1], slots, v, live, layout="HND") + assert torch.equal(ours_raw, ref_raw) + + +def test_the_fused_kv_insert_skips_a_negative_slot() -> None: + """A live row whose slot is the -1 sentinel stores nothing (no write before the pool row).""" + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + tokens, tokens_per_block, pages = 4, 32, 8 + g = torch.Generator(device="cuda").manual_seed(7) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + raw = torch.randint( + 0, + 256, + (pages, 3, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.tensor([5, -1, 40, 77], device="cuda", dtype=torch.int32) + + def view(t): + return t.view(torch.float8_e4m3fn)[:, 1].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, True, pos, False) + torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, tokens + ) + plain = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + keep = torch.tensor([0, 2, 3], device="cuda") + ref_raw = raw.clone() + ref = view(ref_raw) + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D)[keep] + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D)[keep] + write_kv_slots(ref[:, 0], slots[keep], k, 3, layout="HND") + write_kv_slots(ref[:, 1], slots[keep], v, 3, layout="HND") + assert torch.equal(ours_raw, ref_raw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..e3de00587e67 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_attention entry: the Eagle3 drafter layer's attention against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp2``): the target's hybrid NVFP4 +``MiniMaxM3KVCacheManagerV2`` with its one-model Eagle3 draft layer, the requests allocated in it, drafter +``TrtllmAttentionMetadata`` over the manager's FP8 draft view (the manager the runtime's draft swap installs) filled +by its ``prepare()``, and the drafter's ``TrtllmAttention`` built as ``M3Eagle3Draft`` builds it (fused NeoX RoPE at +the drafter checkpoint's theta, FP8 KV). Earlier rows go into the draft layer's cache through the in-tree scatter at +the slots the manager assigned; the call rotates and appends the step's own. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. +# The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: +# it takes the rotated Q in E4M3 (rounding the reference's Q to E4M3 takes the RMS gap from ~9 to ~7 ulps) and rounds +# the softmax probabilities before the PV product. Measured on GB300 over these cells (max element / RMS): 13.6 / 8.6 +# (4 x 4 rows), 9.9 / 7.2 (8 x 1), 20.4 / 7.7 (512 context rows). The band leaves about 1.5x; a reference shifted +# by one row is far outside it. +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 + +# Cached tokens per request of each cell; the new rows per request come from the cell. +_CACHED = { + (2, 4): (200, 1000), + (2, 1): (127, 2047), + (1, None): (0,), +} +#: Rows past the live ones in the fused QKV, as a padded step hands them over. +_PAD = 3 + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + +# Non-paged context FMHA on SM107 does not apply RoPE (this fixture's path; the served drafter runs the paged context +# FMHA, which rotates): the context cell is a strict expected failure until the runtime rotates there too. +_UNROTATED_CONTEXT = pytest.mark.xfail( + strict=True, + reason="non-paged context FMHA on SM107 does not apply RoPE (shared runtime)", +) +_CELLS = [ + pytest.param(c, marks=_UNROTATED_CONTEXT) if c.spec.get("context") else c for c in op.CELLS +] + + +@pytest.mark.parametrize("cell", _CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + context = spec.get("context") + cached = _CACHED[spec["reqs"], spec.get("qlen")] + qlens = (context,) if context else (spec["qlen"],) * spec["reqs"] + assert len(cached) == spec["reqs"] + contexts = 1 if context else 0 + rids = list(range(spec["reqs"])) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = rt.make_draft_metadata(manager) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + assert int(metadata.num_tokens) == total + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + certifying(op) + out = op(backend, qkv, metadata) + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{cell.why}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + group = rt.HEADS // rt.KV_HEADS + for h in range(rt.KV_HEADS): + cols = slice(h * group * rt.HEAD_DIM, (h + 1) * group * rt.HEAD_DIM) + element, rms = _ulp_stats(out[:, cols], ref[:, cols]) + print(f"[{cell.why}] KV head {h}: max element {element:.2f} ulp, rms {rms:.3f} ulp") + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + # The call's append: each live row's rotated K and its V at the row's position, E4M3 (unit KV scale), to + # within one E4M3 step of the float64 rotation (the kernel rotates in FP32). + pool = manager.get_buffers(rt.DRAFT, "HND") + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + for h in range(rt.KV_HEADS): + k_cols = q_size + h * rt.HEAD_DIM + v_cols = q_size + (rt.KV_HEADS + h) * rt.HEAD_DIM + want_k = _rope(rows[:, k_cols : k_cols + rt.HEAD_DIM].float(), positions) + want_v = rows[:, v_cols : v_cols + rt.HEAD_DIM].to(torch.float8_e4m3fn) + got_k = pool[slots // rt.PAGE, 0, h, slots % rt.PAGE].to(torch.float64) + got_v = pool[slots // rt.PAGE, 1, h, slots % rt.PAGE] + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), ( + f"K head {h} is appended rotated, in E4M3" + ) + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)), f"V head {h}" + first += n + + # The band discriminates: a reference whose rows are shifted by one token is outside it. + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py new file mode 100644 index 000000000000..d401a9d532e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's context attention on SM107 as the server runs it: paged context FMHA (chunked prefill +and speculative draft tokens set ``use_paged_context_fmha``) over the shared manager's FP8 draft view at its +128-token pages, which SM107 serves through the thop paged-context kernel (no TRTLLM-Gen FMHA there). + +Each case is checked against the exact float64 reference of ``m3_draft_attention`` (``test_modeling_v2_m3_draft_ +attention.py``'s band), with fresh contexts, chunked appends over a cached prefix (page-aligned or not), two contexts +in one step and a mixed step with generation rows beside the context. The appended K / V rows and a shifted-reference +discriminator are checked as there. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp # noqa: E402 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the SM107 thop paged-context path (no TRTLLM-Gen FMHA)" +) + +# test_modeling_v2_m3_draft_attention.py's band (BF16 ulps of each row's largest magnitude / of the relative RMS). +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 +_PAD = 3 + +# (case id, cached tokens per request, new rows per request, context requests first). +_CASES = [ + ("fresh_ctx_512", (0,), (512,), 1), + ("append_1000_unaligned", (1000,), (700,), 1), + ("append_1024_aligned", (1024,), (300,), 1), + ("two_ctx_prefix_384", (0, 384), (257, 129), 2), + ("prefix_ctx_beside_verify", (500, 200, 1000), (200, 4, 4), 1), + ("fresh_ctx_beside_decode", (0, 200, 1000), (200, 1, 1), 1), + ("fresh_ctx_beside_verify", (0, 200, 1000), (200, 4, 4), 1), +] + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _paged_draft_metadata(manager) -> TrtllmAttentionMetadata: + """``rt.make_draft_metadata`` with the server's runtime features (chunked prefill, Eagle3 draft tokens).""" + view = manager.get_draft_kv_cache_view() + metadata = TrtllmAttentionMetadata( + max_num_requests=rt.MAX_BATCH, + max_num_tokens=rt.MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=rt.mapping(), + runtime_features=AttentionRuntimeFeatures( + chunked_prefill=True, has_speculative_draft_tokens=True + ), + ) + assert metadata.use_paged_context_fmha + return metadata + + +def _prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + """``rt.prepare`` as the serving runner prepares a step: a context request's prompt length is its CURRENT CHUNK's + token count (``prompt_lengths.append(len(prompt_tokens))``) with the prefix in ``num_cached_tokens_per_seq``, so + the chunk is appended at the cached length; generation rows as ``rt.prepare``.""" + from tensorrt_llm._torch.metadata import KVCacheParams + + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +@pytest.mark.parametrize("case", _CASES, ids=[c[0] for c in _CASES]) +def test_paged_context_at_128_token_pages(case) -> None: + name, cached, qlens, contexts = case + rids = list(range(len(qlens))) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(name.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = _paged_draft_metadata(manager) + assert metadata.tokens_per_block == rt.PAGE == 128 + assert 128 in manager.get_draft_kv_cache_view().trtllm_gen_extra_tokens_per_block + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + pool = manager.get_buffers(rt.DRAFT, "HND") + prefix_slots = { + rid: rt.flat_slots(manager, rt.DRAFT, rid, torch.arange(c, device="cuda")).long() + for rid, c in zip(rids, cached) + if c + } + prefix_before = { + rid: pool[sl // rt.PAGE, :, :, sl % rt.PAGE].clone() for rid, sl in prefix_slots.items() + } + _prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + out = op(backend, qkv, metadata) + # The cached prefix's K / V rows are left as they were (the call appends only the step's rows). + torch.cuda.synchronize() + changed = { + rid: int( + ( + pool[sl // rt.PAGE, :, :, sl % rt.PAGE].view(torch.uint8) + != prefix_before[rid].view(torch.uint8) + ) + .any(dim=-1) + .sum() + ) + for rid, sl in prefix_slots.items() + } + print(f"\n[{name}] cached prefix rows x heads x K/V changed by the call: {changed}") + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{name}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + # Every live row's K / V landed in its page and slot (the append of a chunk past a cached prefix). + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + for h in range(rt.KV_HEADS): + v_cols = q_size + (rt.KV_HEADS + h) * rt.HEAD_DIM + want_v = rows[:, v_cols : v_cols + rt.HEAD_DIM].to(torch.float8_e4m3fn) + got_v = pool[slots // rt.PAGE, 1, h, slots % rt.PAGE] + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)), f"V head {h}" + first += n + + assert not any(changed.values()), f"the call overwrote cached prefix K / V rows: {changed}" + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py new file mode 100644 index 000000000000..069e4fdc3063 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for kernels/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V +append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row +first draft pass and the 1-row later passes of 1-2 requests. + +Lengths sweep the KV tile edges each query row in turn ends on (the rows of a Q CTA with no key in its trailing +splits), with mixed per-request lengths; every call is held to the fp32 reference over the cache after the append +(E4M3 query, unit KV scales, no RoPE) within the FP8-KV band, finite, for both split-KV reducers. +""" + +import math +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) + +assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +FP8_TOL = 6e-2 +# Persistent like the FMHA manager's workspace and the spec-decoding metadata of a real step: the in-tree preprocess +# kernels must not see these move between calls. +_PERSISTENT = {} + + +def _persistent(name: str, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _edge_lens(plan, qlen: int): + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +def _library(qlen: int, nreq: int, own_reducer: bool): + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = own_reducer + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = torch.cuda.get_device_properties(0).multi_processor_count + lib._on_taken = lambda rows: None + # The in-tree preprocess pass (its RoPE-free cache append); the one-launch preprocess has its own bitwise test. + lib._fused_rope_mode = None + lib._pdl = False + lib._on_rope_taken = None + return lib + + +def _attn(qlen: int): + return SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=0, theta=10000.0, scale_type=0, scale=1.0, max_positions=1024 + ), + ) + + +def _run(lib, qlen: int, kv_lens, seed: int, spec: bool = True): + """One drafter generation pass over a fresh FP8 cache; returns (output, fp32 reference). ``spec`` passes the + spec-decoding lengths/positions (a multi-token generation step); without them a multi-row pass is the drafter's + plain generation step of input_seq_length rows, as it runs in-model.""" + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g).clamp_(-4, 4) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + for r, n in enumerate( + kv_lens + ): # shift the tail third of V, so a missed trailing split fails the band + start = 2 * n // 3 + for pg in range(start // PAGE, math.ceil(n / PAGE)): + pool[int(perm[r, pg]), 1, :, (start % PAGE if pg == start // PAGE else 0) :] += 1.0 + pool = pool.to(torch.float8_e4m3fn) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + pointers = torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64) + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=dev, generator=g) + seq = torch.tensor(kv_lens, dtype=torch.int32, device=dev) + one = torch.ones(1, dtype=torch.float32, device=dev) + multi = qlen > 1 and spec + out = torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=_attn(qlen), + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=pointers, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=seq, + spec_decoding_generation_lengths=( + _persistent( + f"gen_lengths{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=( + _persistent( + f"position_offsets{qlen}x{nreq}", + lambda: ( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous() + ), + ) + if multi + else None + ), + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=out, + ) + lib.run_generation(params) + torch.cuda.synchronize() + # fp32 reference over the cache after the append: E4M3 query, unit scales, causal over the step's rows. + q = qkv[:, : H * D].float().view(nreq, qlen, H, D).to(torch.float8_e4m3fn).float() + flat = pool.float() + ref = torch.empty(nreq, qlen, H, D, device=dev) + for r, n in enumerate(kv_lens): + blocks = perm[r, : math.ceil(n / PAGE)] + k = flat[blocks, 0].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + v = flat[blocks, 1].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + k = k.repeat_interleave(H // KVH, dim=0) + v = v.repeat_interleave(H // KVH, dim=0) + s = torch.einsum("qhd,hkd->hqk", q[r], k) * D**-0.5 + visible = ( + torch.arange(n, device=dev)[None, :] + < (n - qlen + 1 + torch.arange(qlen, device=dev))[:, None] + ) + s = s.masked_fill(~visible[None], float("-inf")) + ref[r] = torch.einsum("hqk,hkd->qhd", torch.softmax(s, dim=-1), v) + return out.float().view(nreq, qlen, H, D), ref + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_the_drafter_pass_at_tile_edges(qlen, nreq) -> None: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if fullctx.has_own_reducer() else (False,) + for own in reducers: + fullctx.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_LEN, own, get_env_enable_pdl() + ) + lens = _edge_lens(fullctx.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + lens = [n for n in lens if n > qlen] + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for i, kv_lens in enumerate(batches): + # The first draft pass runs in-model as a plain multi-row generation step; every other batch also takes the + # spec-decoding form. + spec = qlen > 1 and i % 2 == 1 + for own in reducers: + got, ref = _run(_library(qlen, nreq, own), qlen, kv_lens, sum(kv_lens), spec=spec) + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +def _drafter_model(scale: float): + """A drafter whose one attention module holds a TrtllmAttention backend as ``attn`` (the backend is no module).""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + backend = TrtllmAttention.__new__(TrtllmAttention) + backend.kv_cache_scaling_factor = torch.full((1,), scale, device="cuda") + backend.is_mla_enable, backend.kv_lora_rank, backend.v_head_dim, backend.head_dim = ( + False, + None, + None, + D, + ) + backend._fmha_manager = SimpleNamespace(fmha_libs=[]) + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.zeros(1, device="cuda")) + model.self_attn = torch.nn.Module() + model.self_attn.attn = backend + return model, backend + + +def test_install_declines_a_scaled_kv_cache() -> None: + """At unit KV scales the arithmetic is the stock decode's; a scaled cache keeps the stock path.""" + reasons = [] + model, layer = _drafter_model(0.5) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 0 + assert layer._fmha_manager.fmha_libs == [] and reasons == ["a scaled FP8 KV cache"] + model, layer = _drafter_model(1.0) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 1 + assert isinstance(layer._fmha_manager.fmha_libs[0], drafter.M3DrafterFullCtxFmha) + + +def test_a_cache_bound_past_the_plan_capacity_keeps_the_stock_path() -> None: + """The plans are compiled for ``max_kv_len`` of K/V: a step whose cache length bound is longer falls through.""" + + from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase + + class _Backend: + num_heads, num_kv_heads, head_dim, quant_mode = H, KVH, D, FP8_KV_QUANT_MODE + kv_scale_orig_quant = kv_scale_quant_orig = None + + backend = _Backend() + backend.kv_scale_orig_quant, backend.kv_scale_quant_orig = object(), object() + reasons = [] + lib = _library(1, 1, False) + lib._on_uncovered = reasons.append + lib._attn_ref = weakref.ref(backend) + fwd = SimpleNamespace( + attention_input_type=None, + kv_scale_orig_quant=backend.kv_scale_orig_quant, + kv_scale_quant_orig=backend.kv_scale_quant_orig, + ) + meta = SimpleNamespace(num_contexts=0, num_generations=1, max_seq_len=MAX_LEN + 8) + q = torch.empty(1, H * D, device="cuda") + assert not lib._is_supported(q, None, None, meta, fwd, phase=FmhaPhase.GENERATION) + assert reasons == ["the cache length bound is beyond the plans' K/V capacity"] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py new file mode 100644 index 000000000000..a27ede466e03 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -0,0 +1,406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for minimax_m3_sm107tp2_drafter_rope_fp8_kv, the TP2 drafter layer's one-launch generation +preprocess: bitwise against the in-tree ``thop.trtllm_gen_generation_preprocess`` followed by the torch E4M3 cast +of its BF16 query, on the query bytes and on every byte of the paged cache, with GPT-NeoX RoPE from a cos/sin table +(full and partial rotary width), multi-token (spec-decoding positions) and plain passes, one and two requests, and +query values past the E4M3 range. The library's fused pass is also held bitwise to its preprocess pass end to end. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp2_drafter_rope_fp8_kv"): + pytest.skip("minimax_m3_sm107tp2_drafter_rope_fp8_kv is not built", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +ROPE_GPT_NEOX = 2 +_PERSISTENT = {} + + +def _persistent(name, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _case(qlen, kv_lens, rotary_dim, multi, seed): + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = (torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g) * 2).to( + torch.float8_e4m3fn + ) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, device=dev, generator=g) * 4 + # A few query / key values near and past the E4M3 range (both casts saturate at 448). + hot = torch.rand(qkv.shape, device=dev, generator=g) < 0.01 + qkv = torch.where( + hot, qkv.sign() * (430 + 80 * torch.rand(qkv.shape, device=dev, generator=g)), qkv + ) + qkv = qkv.to(torch.bfloat16).contiguous() + max_pos = max(kv_lens) + 64 + cos_sin = torch.randn(max_pos, rotary_dim // 2, 2, device=dev, generator=g).float().contiguous() + attn = SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=ROPE_GPT_NEOX, + rotary_inv_freq=None, + rotary_cos_sin=cos_sin, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=rotary_dim, theta=10000.0, scale_type=0, scale=1.0, max_positions=max_pos + ), + ) + one = torch.ones(1, dtype=torch.float32, device=dev) + spec_offsets = None + if multi: + spec_offsets = _persistent( + f"pos{qlen}x{nreq}", + lambda: torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous(), + ) + # Spec positions are not the plain row index: shift them per request. + spec_offsets.copy_( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1) + + torch.arange(nreq, dtype=torch.int32, device=dev)[:, None] + ) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=attn, + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=torch.tensor(kv_lens, dtype=torch.int32, device=dev), + spec_decoding_generation_lengths=( + _persistent( + f"len{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=spec_offsets, + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev), + ) + return params, pool + + +def _preprocess(params, mpc): + """run_generation's in-tree preprocess call and the E4M3 cast of its query.""" + attn, meta, fwd, rope = params.attn, params.meta, params.fwd, params.attn.rope_params + q_processed = thop.trtllm_gen_generation_preprocess( + params.qkv_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope.dim, + rope.theta, + int(rope.scale_type), + rope.scale, + rope.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + False, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + mpc, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + )[0] + rows = params.num_tokens * attn.num_heads * attn.head_dim + return ( + q_processed.flatten()[:rows] + .to(torch.float8_e4m3fn) + .view(params.num_tokens, attn.num_heads, attn.head_dim) + ) + + +def _op(params, pool, mode): + attn = params.attn + nreq, qlen = params.batch_size, params.input_seq_length + kv_pages = pool.view(-1, KVH, PAGE, D) + tables = params.meta.kv_cache_block_offsets[0, :nreq] + q = torch.empty( + params.num_tokens, + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=pool.device, + ) + multi = params.spec_decoding_generation_lengths is not None and qlen > 1 + torch.ops.trtllm.minimax_m3_sm107tp2_drafter_rope_fp8_kv( + params.qkv_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + attn.rope_params.dim, + PAGE, + mode, + True, + ) + return q + + +# One-request multi-token passes keep the in-tree preprocess (it reads an unbuilt token map), so the spec cases hold +# two requests. +CASES = [ + (qlen, kv_lens, rotary, multi) + for qlen, multi in ((1, False), (4, True), (4, False)) + for kv_lens in (((4, 130) if multi else (4,)) if qlen > 1 else (130,), (8191, 131073)) + for rotary in (128, 64) +] + + +def _bytes(t): + return t.contiguous().view(torch.uint8) + + +def _matches(case, seed): + """Per rope mode: (query bytes equal, cache bytes equal) against the in-tree preprocess + cast.""" + qlen, kv_lens, rotary, multi = case + kv_lens = [max(n, qlen) for n in kv_lens] + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(qlen, kv_lens, rotary, multi, seed) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + pool_ref = pool.clone() + out = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + out[mode] = ( + bool(torch.equal(_bytes(q), _bytes(q_ref))), + bool(torch.equal(_bytes(pool), _bytes(pool_ref))), + ) + return out + + +@pytest.mark.parametrize( + "case", + CASES, + ids=lambda c: f"q{c[0]}-{'x'.join(map(str, c[1]))}-r{c[2]}-{'spec' if c[3] else 'plain'}", +) +def test_the_pinned_rope_mode_is_bitwise(case) -> None: + result = _matches(case, seed=7) + print(f"[rope modes] {case}: {result}", flush=True) + assert result[modeling.DRAFTER_ROPE_MODE] == (True, True), result + + +def test_the_fused_pass_matches_the_preprocess_pass() -> None: + mpc = torch.cuda.get_device_properties(0).multi_processor_count + for qlen, nreq in sorted(modeling.FULLCTX_DECODE_RECIPES): + for multi in (False, True) if qlen > 1 else (False,): + kv_lens = [70001 + 517 * r for r in range(nreq)] + outs = [] + for fused in (False, True): + params, pool = _case(qlen, kv_lens, 128, multi, seed=11) + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = True + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = mpc + lib._fused_rope_mode = modeling.DRAFTER_ROPE_MODE if fused else None + lib._pdl = True + lib._one_request_multi = drafter.ROPE_ONE_REQUEST_MULTI + taken = [] + lib._on_rope_taken = lambda rows, what: taken.append(what) + lib._on_taken = lambda rows: None + fullctx.prebuild(torch.device("cuda"), lib._recipes, MAX_LEN, True, True) + lib.run_generation(params) + torch.cuda.synchronize() + assert bool(taken) == ( + fused and (lib._one_request_multi or not (multi and nreq == 1)) + ) + outs.append((params.output.clone(), pool.clone())) + assert torch.equal(_bytes(outs[0][1]), _bytes(outs[1][1])), (qlen, nreq, multi, "cache") + assert torch.equal(_bytes(outs[0][0]), _bytes(outs[1][0])), ( + qlen, + nreq, + multi, + "output", + ) + + +def _round_bf16(x): + u = x.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000 + return u.to(torch.int64).to(torch.int32).view(torch.float32) + + +def _discriminating(n, seed): + """(cos, sin, x, pair) quadruples whose rotated first-half output's E4M3 byte differs between the rope modes + (fp32 CPU emulation; the GPU comparison is the verdict).""" + g = torch.Generator().manual_seed(seed) + keep = [] + while sum(k.shape[0] for k in keep) < n: + m = 1 << 20 + c, s = torch.randn(m, generator=g), torch.randn(m, generator=g) + x = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + xp = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + sn = -s + r0 = (c * x) + (sn * xp) + r1 = ((c.double() * x.double()) + (sn * xp).double()).float() + r2 = ((sn.double() * xp.double()) + (c * x).double()).float() + b = [ + _round_bf16(r.contiguous()).to(torch.float8_e4m3fn).view(torch.uint8) + for r in (r0, r1, r2) + ] + diff = (b[0] != b[1]) | (b[0] != b[2]) | (b[1] != b[2]) + diff &= (r0.abs() < 400) & (r1.abs() < 400) + keep.append(torch.stack([c, s, x, xp], 1)[diff]) + return torch.cat(keep)[:n] + + +def test_crafted_inputs_pin_the_rope_mode() -> None: + """Rotations whose E4M3 byte depends on how the two products are summed: the in-tree preprocess and exactly the + pinned mode agree on every one.""" + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(1, [130], 128, False, seed=3) + half = 64 + quads = _discriminating(H * half, seed=5).cuda() + pos = 129 # the one row's rope position: cache length 130, one query row + table = params.attn.rotary_cos_sin.view(-1, half, 2) + # One (cos, sin) per dim: every head shares the position's table row, so the heads take the same dims' quads with + # their own (x, pair). + table[pos, :, 0] = quads[:half, 0] + table[pos, :, 1] = quads[:half, 1] + qkv = params.qkv_input.view(-1, H + 2 * KVH, D) + for h in range(H): + qkv[0, h, :half] = quads[:half, 2].to(torch.bfloat16) + qkv[0, h, half:] = quads[:half, 3].to(torch.bfloat16) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + agree = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + agree[mode] = int((_bytes(q)[0, :, :half] == _bytes(q_ref)[0, :, :half]).sum()) + print( + f"[rope modes crafted] agreeing first-half query bytes of {H * half}: {agree}", flush=True + ) + assert agree[modeling.DRAFTER_ROPE_MODE] == H * half, agree + + +def test_one_request_multi_token_matches_the_guarded_preprocess() -> None: + """The drafter's first pass at one request (several tokens): the in-tree preprocess places those rows by a + workspace token map it does not build for one request, which the TP2 token-map guard (kernels/m3_tokmap_guard) + writes right before the call. After POISONING the map with a two-request call, the fused op must be bitwise the + guarded preprocess + cast; this is what drafter.ROPE_ONE_REQUEST_MULTI rests on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_tokmap_guard as guard, + ) + + guard.install(torch.device("cuda"), 64, lambda rows: None) + mpc = torch.cuda.get_device_properties(0).multi_processor_count + results = {} + for kv in (130, 8191, 131073, 165521): + poison, _ = _case(4, [131072, 4000], 128, True, seed=11) + _preprocess(poison, mpc) + params, pool = _case(4, [kv], 128, True, seed=5) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + pool_ref = pool.clone() + pool.copy_(start) + q = _op(params, pool, modeling.DRAFTER_ROPE_MODE) + torch.cuda.synchronize() + results[kv] = ( + bool(torch.equal(_bytes(q), _bytes(q_ref))), + bool(torch.equal(_bytes(pool), _bytes(pool_ref))), + ) + print(f"[one-request multi-token vs guarded preprocess] {results}", flush=True) + if drafter.ROPE_ONE_REQUEST_MULTI: + assert all(v == (True, True) for v in results.values()), results diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..918174c1295f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp8_indexer_qk_norm_rope entry of the MiniMax-M3 tp2 target: index-Q out and index-K in the +cache against the in-place fused norm + RoPE followed by the E4M3 cast, byte for byte up to the 0.1% budget the +in-tree test of the kernel allows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 + m3_fp8_indexer_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" + +DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 + + +def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: + assert a.shape == b.shape and a.dtype == b.dtype == torch.float8_e4m3fn + differ = int((a.view(torch.uint8) != b.view(torch.uint8)).sum()) + assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + m, heads = cell.spec["m"], cell.spec["heads"] + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + qk = torch.randn(m, (heads + 1) * DIM, generator=g, device="cuda").to(torch.bfloat16) + qw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (m,), generator=g, device="cuda", dtype=torch.int32) + pages = (m + PAGE - 1) // PAGE + 1 + cache = torch.zeros(pages, 1, PAGE, DIM, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(pages * PAGE, generator=g, device="cuda")[:m].to(torch.int32) + before = qk.clone() + args = (qk, cache, slots, heads, DIM, ROTARY, EPS, qw, kw, THETA, pos) + with validating(op): + q = op(*args) + assert torch.equal(qk, before), "the BF16 input is left untouched" + q_ref, k_ref = op.reference(*args) + _bytes_close(q.reshape(m, heads, DIM), q_ref) + flat = cache.view(-1, DIM) + _bytes_close(flat[slots.long()], k_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py new file mode 100644 index 000000000000..f58dc5f3f0ce --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""paged_cache.fp8_e4m3_saturate (FP8_SATURATE): the E4M3 cache writes and the transient FP8 query casts saturate to ++-448 like the CUDA writers (__NV_SATFINITE); every finite in-range value keeps the plain cast's bits, values already +in E4M3 pass through, and the kill switch is exactly the plain cast. Where torch's cast already saturates (this image) +the conversion is the cast alone, one kernel; otherwise it clamps first, the same bytes.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import paged_cache + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fp8_e4m3_saturate_matches_satfinite(dtype, monkeypatch) -> None: + g = torch.Generator(device="cuda").manual_seed(0) + inside = (torch.randn(4096, generator=g, device="cuda") * 100).clamp(-448, 448).to(dtype) + plain = inside.to(torch.float8_e4m3fn) + assert torch.equal( + paged_cache.fp8_e4m3_saturate(inside).view(torch.uint8), plain.view(torch.uint8) + ) + edge = torch.tensor( + [449.0, 463.0, 464.0, 1e4, float("inf"), -500.0, -float("inf")], dtype=dtype, device="cuda" + ) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0], device="cuda") + assert torch.equal(paged_cache.fp8_e4m3_saturate(edge).float(), want) + nan = paged_cache.fp8_e4m3_saturate(torch.tensor([float("nan")], dtype=dtype, device="cuda")) + assert torch.isnan(nan.float()).all() + already = edge.clamp(-448, 448).to(torch.float8_e4m3fn) + assert paged_cache.fp8_e4m3_saturate(already) is already + monkeypatch.setattr(paged_cache, "FP8_SATURATE", False) + killed = paged_cache.fp8_e4m3_saturate(edge) + assert torch.equal(killed.view(torch.uint8), edge.to(torch.float8_e4m3fn).view(torch.uint8)) + + +def test_cache_write_saturates() -> None: + """The cache-dtype conversion every write_kv_slots branch applies (compared as bytes: torch has no float8 + index_copy_ kernel), and an E4M3 source passing through untouched.""" + values = torch.tensor([[1.0, 500.0, -1e5, float("inf")]], dtype=torch.bfloat16, device="cuda") + out = paged_cache._to_cache_dtype(values, torch.float8_e4m3fn) + assert out.dtype == torch.float8_e4m3fn + want = torch.tensor([[1.0, 448.0, -448.0, 448.0]], device="cuda").to(torch.float8_e4m3fn) + assert torch.equal(out.view(torch.uint8), want.view(torch.uint8)) + assert paged_cache._to_cache_dtype(out, torch.float8_e4m3fn) is out + bf = paged_cache._to_cache_dtype(values, torch.bfloat16) + assert bf.dtype == torch.bfloat16 and torch.equal(bf, values) + + +_EDGES = [ + 448.0, + 449.0, + 463.0, + 464.0, + 1e4, + float("inf"), + -448.0, + -449.0, + -464.0, + -1e4, + -float("inf"), + float("nan"), + 0.0, +] + + +def _cuda_kernels(fn) -> int: + from torch.profiler import ProfilerActivity, profile + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + return sum(1 for e in prof.events() if e.device_type.name == "CUDA") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("cast_saturates", [True, False], ids=["cast_only", "clamp_fallback"]) +def test_fp8_e4m3_saturate_bytes_equal_clamp_then_cast(dtype, cast_saturates, monkeypatch) -> None: + """Both paths give the bytes of clamp(+-448) followed by the cast, edge values included.""" + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", cast_saturates) + x = torch.tensor(_EDGES, dtype=dtype, device="cuda") + want = x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + assert torch.equal(paged_cache.fp8_e4m3_saturate(x).view(torch.uint8), want.view(torch.uint8)) + + +def test_fp8_e4m3_saturate_takes_the_cast_only_path_here(monkeypatch) -> None: + """This image's torch cast saturates, on the CPU (the import-time probe) and on CUDA, so the conversion is one + kernel; with the probe failing it falls back to clamp + cast, two kernels.""" + assert paged_cache._CAST_SATURATES + x = torch.tensor(_EDGES, dtype=torch.bfloat16, device="cuda") + assert torch.equal( + x.to(torch.float8_e4m3fn).view(torch.uint8), + x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).view(torch.uint8), + ) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 1 + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", False) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 2 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py new file mode 100644 index 000000000000..cb20a22a51d0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for the attention/m3_fullctx_decode entry: the dense layers' full-context FP8 decode on PrimTS +split-KV decode at the recipe the target pins per (query rows per request, requests). + +Every cell runs against the fp32 reference (q rounded to E4M3 first, as the kernel runs it), eager and replayed from a +CUDA graph (the same bits), and twice (the same bits). The K/V pages are read through the strided views of a flat +sub-page pool, as the KV cache manager hands them out. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_fullctx_decode import ( # noqa: E501 + HEAD_DIM, + HEADS, + KV_HEADS, + MAX_KV_LEN, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_fullctx_decode import ( # noqa: E501 + m3_fullctx_decode as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) + +PAGE = 128 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the entry asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor, subpages: int): + self.pool = pool + self.subpages = subpages + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, self.subpages + + +def _inputs(kv_lens, qlen: int, seed: int, subpages: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * subpages, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1, so a split that missed the tail fails the band. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + for page in range(start // PAGE, (n + PAGE - 1) // PAGE): + v = int(block_table[b, page]) * subpages + 1 + pool[v, :, (start % PAGE if page == start // PAGE else 0) :] += 1.0 + # q as the layer hands it over on decode steps: the FP8 query columns of the fused norm/RoPE kernel's + # [q | k | v] output (row stride wider than q), so no cast makes it compact. + qkv = torch.randn( + (batch * qlen, (HEADS + 2 * KV_HEADS) * HEAD_DIM), device="cuda", generator=g + ).to(torch.float8_e4m3fn) + q = qkv[:, : HEADS * HEAD_DIM].view(batch * qlen, HEADS, HEAD_DIM) + seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32) + return q, _Pool(pool.to(torch.float8_e4m3fn), subpages), block_table, seq_lens + + +@pytest.mark.parametrize("subpages", [2, 3], ids=["kv", "kv+index"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, subpages, certifying) -> None: # noqa: F811 + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], qlen, zlib.crc32(cell.why.encode()) % 2**31, subpages + ) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + certifying(op) + ref = op.reference(*args) + eager = op(*args).clone() + torch.testing.assert_close(eager.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + again = op(*args).clone() + assert torch.equal(again, eager), "not deterministic" + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args) + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager), "the CUDA-graph replay differs from the eager call" + + +def test_the_cells_cover_the_target_rows() -> None: + certified = {(c.spec["qlen"], len(c.spec["kv_lens"])) for c in op.CELLS} + assert set(modeling.FULLCTX_DECODE.keys) <= certified + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: + """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once + in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + + if not kernel.has_own_reducer(): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_fullctx_reduce") + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs(spec["kv_lens"], qlen, 11, 2) + outs = [] + for own in (True, False): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + op(q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe, own_reducer=own) + outs.append(out.float()) + ours, primts = outs + _, exp = torch.frexp(primts.abs()) + spacing = torch.where( + primts == 0, + torch.full_like(primts, 2.0**-133), + torch.ldexp(torch.ones_like(primts), exp - 8), + ) + assert bool(((ours - primts).abs() <= spacing).all()) + + +def _tile_edge_kv_lens(plan, qlen: int): + """KV lengths around every kind of boundary the split ranges have, so each query row of the Q CTA in turn ends + its keys on a KV tile edge (k * tile + d, d in -1 .. qlen - 1): the first tiles, the reducer's 8-group stride, the + lengths where a split's range grows (splits * insts tiles) and a split edge past it.""" + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: + """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys + (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode + step first hit (NaN rows from exp2(-inf - -inf)).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, False, get_env_enable_pdl() + ) + lens = _tile_edge_kv_lens(kernel.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, sum(kv_lens), 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: + """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just + past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + lengths = [(1 << 20) + 1, (1 << 20) + 8] + ([1_500_001] if nreq == 1 else []) + batches = [[n] * nreq if nreq == 1 else [n, (1 << 20) - 3] for n in lengths] + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, 13, 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, own, get_env_enable_pdl() + ) + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_identity_selection.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_identity_selection.py new file mode 100644 index 000000000000..0c2af37bf5ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_identity_selection.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the short-sequence identity selection (runtime policy ``short_sequence_identity``) is the stock table. + +On a step whose every query token has at most 16 valid blocks, this target skips the index scorer and the select and +stages the table (i < n_valid ? i : -1). Stock runs the in-tree CuTe DSL scorer and ``select_blocks_from_maxscore``, +which returns ascending ids and drops a valid block only when its max score is -inf. The scorer's scores are FP8 +products summed in FP32 (|score| <= 448^2 * 128), so -inf comes only from the causal mask, which n_valid excludes. +This test drives the real scorer with extreme FP8 operands at the TP2 geometry and checks both claims. +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) +from tensorrt_llm._torch.custom_ops import ( + cute_dsl_custom_ops, # noqa: F401 — registers the scorer op +) + +PAGE, DIM, TOPK, HEADS, QLEN = 128, 128, 16, 2, 4 +E4M3_MAX = 448.0 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + + +def _fill(kind, shape, gen): + if kind == "random": + x = torch.randn(shape, generator=gen, device="cuda") * 64 + elif kind == "max": + x = torch.full(shape, E4M3_MAX, device="cuda") + elif kind == "negmax": + x = torch.full(shape, -E4M3_MAX, device="cuda") + elif kind == "alternating": + x = torch.full(shape, E4M3_MAX, device="cuda") + x[..., 1::2] = -E4M3_MAX + elif kind == "zero": + x = torch.zeros(shape, device="cuda") + else: # the smallest E4M3 subnormal + x = torch.full(shape, 2.0**-9, device="cuda") + return x.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + + +_KINDS = ("random", "max", "negmax", "alternating", "zero", "tiny") + + +@pytest.mark.parametrize( + "seq_lens", [(4,), (129,), (2048,), (700, 2048), (5, 1500, 2047, 300), (1, 128, 256, 1024)] +) +@pytest.mark.parametrize("q_kind", _KINDS) +@pytest.mark.parametrize("k_kind", _KINDS) +def test_stock_selection_of_a_short_step_is_the_identity_table(seq_lens, q_kind, k_kind) -> None: + gen = torch.Generator(device="cuda").manual_seed(0) + reqs, total_q = len(seq_lens), len(seq_lens) * QLEN + nblk = [(s + PAGE - 1) // PAGE for s in seq_lens] + pages = sum(nblk) + 4 + perm = torch.randperm(pages, generator=gen, device="cuda").int() + block_table = torch.zeros((reqs, max(nblk)), dtype=torch.int32, device="cuda") + start = 0 + for r, n in enumerate(nblk): + block_table[r, :n] = perm[start : start + n] + start += n + n_valid = torch.tensor( + [(max(s - QLEN + o + 1, 0) + PAGE - 1) // PAGE for s in seq_lens for o in range(QLEN)], + dtype=torch.int32, + device="cuda", + ) + idx_q = _fill(q_kind, (total_q, HEADS, DIM), gen) + idx_k = _fill(k_kind, (pages, PAGE, DIM), gen) + max_score = torch.full((HEADS, max(nblk), total_q), float("-inf"), device="cuda") + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, + idx_k, + block_table, + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + max_score.transpose(1, 2), + QLEN, + ) + for tok, nv in enumerate(n_valid.tolist()): + assert bool(torch.isfinite(max_score[:, :nv, tok]).all()), ( + f"token {tok}: a causally valid block scored non-finite" + ) + table = select_blocks_from_maxscore( + max_score, + topk=TOPK, + n_valid_blocks=n_valid, + init_blocks=1, + local_blocks=4, + head_major_output=True, + ).to(torch.int32) + # The selector returns [tokens, kv_heads, topk]; the decode reads it as [kv_heads, tokens, topk] (layers.py). + table = table.permute(1, 0, 2) + i = torch.arange(TOPK, dtype=torch.int32, device="cuda") + identity = torch.where( + i[None, :] < n_valid[:, None], i[None, :], torch.full_like(i[None, :], -1) + ) + assert torch.equal(table, identity[None].expand(HEADS, total_q, TOPK)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..85c5e6b746c1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score entry: the MSA index scorer with its pre-wait schedule. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score`` at its own 256-CTA split) bit for bit, which the +pre-wait prefetch, the early tile issue and the cell's split must not change. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_index_score import ( # noqa: E501 + PAGE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 2 # index heads per rank at TP2 +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + assert spec["heads"] == HEADS + args = (idx_q, index_k, block_table, seq, score, qlen, spec["split_k"], spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, 1536, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4 + ) + assert torch.equal(score, stock) + + +@pytest.mark.parametrize("qlen", [1, 2, 3, 4]) +def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> None: + """Lengths around the edge between pre-staged (all-old) and post-wait first blocks, 1 and 2 requests.""" + for k in (1, 2, 1294): + for length in (k * PAGE - qlen - 1, k * PAGE - qlen, k * PAGE - qlen + 1, k * PAGE + 1): + for seq_lens, split in (((length,), 1536), ((length, max(qlen, length // 3)), 768)): + idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) + score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) + op(idx_q, index_k, block_table, seq, score, 4, split, True) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + assert torch.equal(score, stock), (qlen, seq_lens) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, 1536, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..4f7e66dd8910 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_msa_attention entry: one layer's MSA backend step against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp2``): a hybrid NVFP4 ``MiniMaxM3KVCacheManagerV2`` +with the requests allocated in it, ``MiniMaxM3MsaSparseAttentionMetadata`` filled by its ``prepare()``, and the layer's +``MiniMaxM3MsaSparseAttention`` built as the target builds it. Earlier rows (and, on a sparse layer, the step's own +rows, which the target's producer writes before the call) go into the cache through the in-tree scatters, at the +slots the manager assigned. The sparse layers' index-K rows are planted so every block's score is 1/8 or more apart +from every other's: the top-k selection the reference recomputes is then unambiguous. + +The cells are the target's calls (``_attention_core`` in modeling.py): a dense layer with live bf16 K / V in a +prefill chunk and live FP8 K / V in a decode step (the dense layers sit in the hybrid cache's FP8 half), a sparse +layer whose producer already wrote the caches (``k = v = None``) in a prefill chunk beyond 16 blocks and in a mixed +step, both with the NVFP4 scales the target passes. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_msa_attention import ( # noqa: E501 + m3_msa_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" + +# Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, +# per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 +# FMHA does), so what is left is each FMHA's own arithmetic, measured on GB300 over these cells (max element / RMS): +# * fmha_sm100 over the FP8 pool (dense context rows): 10.1 / 6.1. It runs E4M3 x E4M3 and its softmax +# probabilities are rounded before the PV product. +# * trtllm-gen over the FP8 pool (dense generation rows): 7.0 / 6.1. It rounds the probabilities to E4M3 (scaled by +# 448 against the running maximum): emulating that brings a single-tile request to 0.6 RMS. +# * the NVFP4 CSR prefill (sparse context rows): 34.4 / 11.8 at 4096 rows, 26.8 / 9.1 at 512. It stages the +# reconstructed K / V in E4M3 (the in-tree A/B test, test_minimax_m3_sparse_attn_decode, notes it and holds it +# to 3e-2). +# * the Triton NVFP4 sparse decode (sparse generation rows): 1.0 / 0.25 -- the band of its own entry +# (m3_sparse_decode): exact operands, FP16 probabilities, BF16 rounding. +# Each band leaves about 1.5x the measured maximum; a shifted or wrongly selected reference is far outside them. +_FP8_CONTEXT = dict(element_ulp=16.0, rms_ulp=9.0) +_FP8_DECODE = dict(element_ulp=12.0, rms_ulp=9.0) +_NVFP4_CONTEXT = dict(element_ulp=48.0, rms_ulp=16.0) +_NVFP4_DECODE = dict(element_ulp=2.0, rms_ulp=0.5) +_BANDS = { + ("dense", "context"): (_FP8_CONTEXT, None), + ("dense", "generation"): (None, _FP8_DECODE), + ("sparse", "context"): (_NVFP4_CONTEXT, None), + ("sparse", "mixed"): (_NVFP4_CONTEXT, _NVFP4_DECODE), +} + + +def _assert_in_bands(output, ref, bands, context_rows: int) -> None: + """Gate the context rows and the generation rows each on their own kernel's band.""" + context_band, decode_band = bands + if context_rows: + assert_within_ulp(output[:context_rows], ref[:context_rows], **context_band) + if context_rows < output.shape[0]: + assert_within_ulp(output[context_rows:], ref[context_rows:], **decode_band) + + +# (request rows, cached tokens per request, new tokens per request, context rows) per cell. +_SHAPES = { + ("dense", "context"): ((384,), (1024,), 1), + ("dense", "generation"): ((128, 2500), (1, 1), 0), + ("sparse", "context"): ((0,), (4096,), 1), + ("sparse", "mixed"): ((0, 2999, 4500), (512, 1, 1), 1), +} + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """The two quantities ``assert_within_ulp`` gates, in BF16 ulps.""" + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _shift(x: torch.Tensor) -> torch.Tensor: + """``x``'s rows moved down by one (byte-wise, so E4M3 rows move too).""" + return x.view(torch.uint8).roll(1, 0).view(x.dtype) if x.element_size() == 1 else x.roll(1, 0) + + +def _shape(cell): + spec = cell.spec + cached, qlens, contexts = _SHAPES[spec["layer"], spec["phase"]] + if spec["phase"] == "context": + assert sum(qlens) == spec["tokens"] + elif spec["phase"] == "generation": + assert len(qlens) == spec["reqs"] + else: + assert sum(qlens[:contexts]) == spec["tokens"] and len(qlens) - contexts == spec["reqs"] + return cached, qlens, contexts + + +def _run_dense(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + rt.write_fp8_rows(manager, rt.DENSE, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + # A prefill chunk arrives in BF16 (the hybrid cache's dense layers take FP8 Q / K / V only on decode steps). + dtype = torch.bfloat16 if contexts else torch.float8_e4m3fn + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + k = torch.randn(total, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + v = torch.randn(total, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + # The NVFP4 target passes its layer scales on every layer; an FP8 layer's write and FMHA take none. + inv = torch.tensor([1.0, 0.5, 2.0], device="cuda") + fwd = 1.0 / inv + return (q, k, v, None, inv, fwd), dict(rids=rids, cached=cached, qlens=qlens, contexts=contexts) + + +def _run_sparse(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + u0, u1 = rt.index_directions(gen) + lengths = [c + n for c, n in zip(cached, qlens)] + keys = [ + torch.randn(n, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths + ] + values = [ + torch.randn(n, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths + ] + quant = rt.nvfp4_quant(torch.cat(keys), torch.cat(values)) + for rid, n, kk, vv in zip(rids, lengths, keys, values): + # Every row, the step's own included: the producer writes a sparse layer's caches before the call. + idx_k = rt.planted_index_keys(n, u0, u1, gen) + rt.write_nvfp4_rows( + manager, + metadata, + rt.SPARSE, + rid, + torch.arange(n), + kk.to(torch.bfloat16), + vv.to(torch.bfloat16), + idx_k, + quant, + ) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + # Query tokens alternate between the two planted directions, so neighbouring tokens select different blocks; the + # rank's index heads take alternating directions, so with more than one KV head they select different blocks too. + turn = torch.arange(total, device="cuda") % 2 + directions = torch.stack([u0, u1]) + idx_q = torch.cat([directions[(turn + h) % 2] for h in range(rt.INDEX_HEADS)], dim=1) + idx_q = idx_q.to(torch.float8_e4m3fn) + return (q, None, None, idx_q, quant, 1.0 / quant), dict( + rids=rids, cached=cached, qlens=qlens, contexts=contexts + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + sparse = spec["layer"] == "sparse" + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=False) + try: + backend = rt.make_msa_backend(rt.SPARSE if sparse else rt.DENSE) + metadata = rt.make_msa_metadata(manager) + run = _run_sparse if sparse else _run_dense + (q, k, v, idx_q, inv, fwd), step = run(cell, gen, manager, backend, metadata) + output = torch.full( + (q.shape[0], rt.HEADS * rt.HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + certifying(op) + result = op(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert result is output + ref = op.reference(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert torch.isfinite(output.float()).all(), "every live row is written" + bands = _BANDS[spec["layer"], spec["phase"]] + context_rows = sum(step["qlens"][: step["contexts"]]) + for name, rows in ( + ("context", slice(0, context_rows)), + ("generation", slice(context_rows, None)), + ): + if output[rows].numel(): + element, rms = _ulp_stats(output[rows], ref[rows]) + print( + f"\n[{cell.why}] {name} rows: max element {element:.2f} ulp, rms {rms:.3f} ulp" + ) + _assert_in_bands(output, ref, bands, context_rows) + + if k is not None: + # The call's own cache write: the live K / V rows land, as a plain E4M3 cast, at their positions. + pool = manager.get_buffers(rt.DENSE, "HND") + first = 0 + for rid, c, n in zip(step["rids"], step["cached"], step["qlens"]): + slots = rt.flat_slots(manager, rt.DENSE, rid, torch.arange(c, c + n)).long() + for role, src in ((0, k), (1, v)): + for h in range(rt.KV_HEADS): + got = pool[slots // rt.PAGE, role, h, slots % rt.PAGE] + want = src[first : first + n, h * rt.HEAD_DIM : (h + 1) * rt.HEAD_DIM] + want = want.to(torch.float8_e4m3fn) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose query rows are shifted by one token is outside it. + shifted = op.reference(backend, _shift(q), k, v, idx_q, metadata, output, sparse, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, shifted, bands, context_rows) + if sparse: + # And the selection is load-bearing: attending every causal block instead is outside it too. + dense = op.reference(backend, q, k, v, idx_q, metadata, output, False, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, dense, bands, context_rows) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py new file mode 100644 index 000000000000..baf8ef7a2b73 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-token valid-block count is written only before the step's first layer. + +The long-row select (``kernels/m3_select_long_dsl.py``, ``SELECT_N_VALID_PREWAIT``) reads its row's entry of +``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the +buffer happens before the step's first layer: + +- the allocation (``_alloc_msa_proxy_scratch``); +- the host-to-device copy in ``prepare()`` (``_build_step_plans``); +- ``on_update_kv_lens``, which the runner calls before the forward. It writes through the plain + ``<<<>>>`` launch of ``msa_kv_update_kernel`` or through torch copies, so any PDL grid of the step starts only + after it completes. + +This test scans ``tensorrt_llm/_torch`` and fails on any other write site (including writes through a local alias of +the buffer), so a new writer has to revisit the pre-wait read. A buffer handed to an opaque kernel launch is not +tracked; such a launch inside a layer must be reviewed by hand. + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +_NAME = "msa_n_valid_blocks" +# In-place tensor methods that write their receiver. +_MUTATORS = { + "copy_", + "fill_", + "zero_", + "index_put_", + "index_copy_", + "scatter_", + "masked_fill_", + "add_", + "sub_", +} +# Calls whose first argument is the destination they write: ``step_inputs.copy_(dst, src)`` and the H2D helpers. +_COPY_HELPERS = {"copy_", "copy_h2d_if_changed", "copy_h2d_pool_shared"} +_BACKEND = "attention/backends/sparse/minimax_m3/msa_backend.py" +# (file relative to tensorrt_llm/_torch, enclosing function) -> writes allowed there +_ALLOWED = { + (_BACKEND, "_alloc_msa_proxy_scratch"), # the allocation + (_BACKEND, "_build_step_plans"), # the H2D copy in prepare() + (_BACKEND, "on_update_kv_lens"), # the device-side correction before the forward +} + + +def _base(node): + """The innermost receiver of an attribute / subscript / call chain.""" + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute) and node.attr == _NAME: + return node + node = node.func if isinstance(node, ast.Call) else node.value + return node + + +def _mentions(node, aliases) -> bool: + """Whether the expression's receiver chain is ``*.msa_n_valid_blocks`` or a local alias of it.""" + b = _base(node) + if isinstance(b, ast.Attribute): + return b.attr == _NAME + return isinstance(b, ast.Name) and (b.id == _NAME or b.id in aliases) + + +def _refers(node, aliases) -> bool: + """Whether any sub-expression (e.g. a branch of a conditional) names the buffer or an alias.""" + return any( + (isinstance(n, ast.Attribute) and n.attr == _NAME) + or (isinstance(n, ast.Name) and (n.id == _NAME or n.id in aliases)) + for n in ast.walk(node) + ) + + +def _write_sites(tree: ast.AST): + sites = [] + + def visit(node, func, aliases): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + func, aliases = node.name, set() + if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + value = node.value + for t in targets: + if isinstance(t, ast.Name) and value is not None and _refers(value, aliases): + aliases.add(t.id) # a local view of the buffer + elif ( + not isinstance(t, ast.Name) + and _mentions(t, aliases) + and not isinstance(node, ast.AnnAssign) + ): + sites.append((func, node.lineno)) + elif ( + isinstance(t, ast.Attribute) + and t.attr == _NAME + and isinstance(node, ast.AnnAssign) + ): + pass # a dataclass field declaration + if isinstance(node, ast.Call): + f = node.func + if isinstance(f, ast.Attribute) and f.attr in _MUTATORS and _mentions(f.value, aliases): + sites.append((func, node.lineno)) + name = ( + f.id + if isinstance(f, ast.Name) + else (f.attr if isinstance(f, ast.Attribute) else "") + ) + if name in _COPY_HELPERS and node.args and _mentions(node.args[0], aliases): + sites.append((func, node.lineno)) + for child in ast.iter_child_nodes(node): + visit(child, func, aliases) + + visit(tree, None, set()) + return sites + + +def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: + found = set() + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if _NAME not in src: + continue + rel = path.relative_to(_ROOT).as_posix() + for func, line in _write_sites(ast.parse(src)): + assert (rel, func) in _ALLOWED, ( + f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " + "wait (kernels/m3_select_long_dsl.py, SELECT_N_VALID_PREWAIT), so a writer inside the " + "step must be ordered before the select's grid or the read moved after the wait" + ) + found.add((rel, func)) + assert found == _ALLOWED, f"expected write sites {_ALLOWED}, found {found}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..dcdfc973c8ba --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_producer entry: the pre-wait-staged NVFP4 QKV + indexer producer at the TP2 heads. + +Every cell is held to the native producer (``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert``) bit for bit on q, +index-q and every K / V / scale / index-K cache byte, and to the fp32 torch reference: the FP8 outputs within one E4M3 +step, the decoded NVFP4 K / V within one E2M1 step of their block scale. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_producer import ( # noqa: E501 + FP4_BLOCK, + HEADS_INDEX, + HEADS_KV, + HEADS_PACKED, + HEADS_Q, + PAGE, + decode_nvfp4_token, + scale_index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_producer import ( # noqa: E501 + m3_producer as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_producer requires a CUDA device" + +PAGES = 1024 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _caches(): + return ( + torch.zeros(PAGES, 2, HEADS_KV, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, HEADS_KV, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _inputs(tokens: int, scale: float, bad_slots: bool, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + weights = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + packed = ( + torch.randn(tokens, HEADS_PACKED * 128, device="cuda", generator=g) * scale + ).bfloat16() + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return packed, slots, quant, weights, rotary, positions + + +def _args(inputs, caches, early: bool = True): + packed, slots, quant, weights, rotary, positions = inputs + return (packed, *caches, slots, quant, EPS, *weights, rotary, positions, HEADS_Q, True, early) + + +def _native(inputs, caches): + packed, slots, quant, weights, rotary, positions = inputs + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, + *caches, + slots, + quant, + HEADS_Q, + HEADS_KV, + HEADS_INDEX, + 128, + 64, + EPS, + *weights, + rotary, + positions, + ) + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + early = spec.get("early", True) + inputs = _inputs( + spec["tokens"], spec["scale"], spec.get("bad_slots", False), zlib.crc32(cell.why.encode()) + ) + stock_caches, caches = _caches(), _caches() + stock = _native(inputs, stock_caches) + args = _args(inputs, caches, early) + certifying(op) + ref = op.reference(*args) + out = op(*args) + torch.cuda.synchronize() + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, s, o in zip(names, (*stock, *stock_caches), (*out, *caches)): + assert torch.equal(s.view(torch.uint8), o.view(torch.uint8)), ( + f"{name} differs from the native producer" + ) + + _within_fp8(out[0], ref["q"], "q") + _within_fp8(out[1], ref["index_q"], "index_q") + kv_data, kv_scale, index_k = caches + slots = inputs[1] + checked = 0 + for t in range(spec["tokens"]): + slot = int(slots[t]) + if slot < 0 or slot >= PAGES * PAGE: + continue + page, within = divmod(slot, PAGE) + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") + for (plane, name), head in ((p, h) for p in ((0, "k"), (1, "v")) for h in range(HEADS_KV)): + scales = kv_scale[page, plane, head].reshape(-1)[scale_index(within, plane).cuda()] + dq = decode_nvfp4_token(kv_data[page, plane, head, within], scales, KV_QUANT[1 + plane]) + r = ref[name][t, head].bfloat16().float() + # Round to nearest E2M1 under the block's E4M3-rounded scale s (in value units s / quant_scale): at most + # half the widest E2M1 step (one unit), plus the clipping at 6 units when s rounded down; a scale that + # underflows E4M3 zeroes the block. A slack of one E4M3 step covers the kernel's approximate reciprocals. + qs = KV_QUANT[1 + plane] + amax = r.view(-1, FP4_BLOCK).abs().amax(dim=1) + unit = (qs * amax / 6.0).to(torch.float8_e4m3fn).float() / qs + bound = torch.where( + unit > 0, 1.15 * unit + (amax - 6.0 * unit / 1.15).clamp_min(0.0), amax + ) + bound = bound.repeat_interleave(FP4_BLOCK) + 1e-12 + assert bool(((dq - r).abs() <= bound).all()), (name, t, head) + checked += 1 + assert checked >= spec["tokens"] - 2 + + +def test_early_trigger_does_not_change_the_bits() -> None: + inputs = _inputs(48, 1.0, False, 7) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=True)) + out_b = op(*_args(inputs, b, early=False)) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +@pytest.mark.parametrize("tokens", [1, 4, 8, 48]) +def test_prewait_step_inputs_do_not_change_the_bits(tokens: int) -> None: + """Reading the position ids, slots and RoPE rows before the grid wait leaves every output and cache byte as is, + padded rows (slot -1) and a slot past the pool included.""" + inputs = _inputs(tokens, 1.0, tokens >= 3, 11) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a), prewait_step_inputs=True) + out_b = op(*_args(inputs, b), prewait_step_inputs=False) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +def test_guard_refuses_a_repeated_slot() -> None: + inputs = list(_inputs(4, 1.0, False, 3)) + inputs[1] = inputs[1].clone() + inputs[1][1] = inputs[1][0] + with pytest.raises(AssertionError, match="must not repeat a slot"): + with validating(op): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_operand_off_the_device() -> None: + inputs = list(_inputs(4, 1.0, False, 9)) + inputs[5] = inputs[5].cpu() # position ids on the host + with pytest.raises(RuntimeError, match="packed's CUDA device"): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_unaligned_index_k_page_stride() -> None: + """The index-K rows are written as 32-bit words: a page stride that is not a multiple of 4 is refused.""" + kv_data, kv_scale, _ = _caches() + backing = torch.zeros(PAGES * (128 * 128 + 1), dtype=torch.float8_e4m3fn, device="cuda") + index_k = backing.as_strided((PAGES, 1, PAGE, 128), (128 * 128 + 1, 128 * 128, 128, 1)) + with pytest.raises(RuntimeError, match="index_k"): + op(*_args(_inputs(4, 1.0, False, 13), (kv_data, kv_scale, index_k))) + + +def test_op_refuses_unaligned_weights_or_rope_table() -> None: + """The norm weights are read as 8-byte words and the RoPE rows as float4: offset views are refused.""" + packed, slots, quant, weights, rotary, positions = _inputs(4, 1.0, False, 17) + offset_weight = torch.zeros(129, dtype=torch.bfloat16, device="cuda")[1:] + with pytest.raises(RuntimeError, match="norm weights"): + op( + *_args( + (packed, slots, quant, [offset_weight, *weights[1:]], rotary, positions), _caches() + ) + ) + flat = torch.zeros(rotary.numel() + 1, device="cuda") + offset_rope = flat[1:].view(rotary.shape) + with pytest.raises(RuntimeError, match="rotary_cos_sin"): + op(*_args((packed, slots, quant, weights, offset_rope, positions), _caches())) + + +def test_each_kv_head_lands_in_its_own_block() -> None: + """K / V of KV head 1 go to head block 1 of the HND pools and differ from head 0's (distinct packed rows).""" + inputs = _inputs(4, 1.0, False, 21) + caches = _caches() + op(*_args(inputs, caches)) + torch.cuda.synchronize() + kv_data = caches[0] + page, within = divmod(int(inputs[1][0]), PAGE) + for plane in (0, 1): + assert not torch.equal(kv_data[page, plane, 0, within], kv_data[page, plane, 1, within]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..bcc0ad67c843 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_qk_norm_rope_fp8 entry of the MiniMax-M3 tp2 target: the out-of-place FP8 form of the fused +per-head Gemma RMSNorm + partial NeoX RoPE against the in-place kernel followed by the E4M3 cast. V is a cast, so it +matches bit for bit; Q and K round once from fp32 where the composition rounds through bf16 first, so each element is +within one E4M3 step of it.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 + m3_qk_norm_rope_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" + +HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 + + +def _args(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q, kv = spec["q"], spec["kv"] + qkv = torch.randn(spec["m"], (q + 2 * kv) * HEAD_DIM, generator=g, device="cuda").to( + torch.bfloat16 + ) + qw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (spec["m"],), generator=g, device="cuda", dtype=torch.int32) + return ( + qkv, + q, + kv, + kv, + HEAD_DIM, + ROTARY_DIM, + EPS, + qw, + kw, + THETA, + True, + pos, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = args[0].clone() + with validating(op): + out = op(*args) + assert torch.equal(args[0], before), "the FP8 form must leave qkv untouched" + ref = op.reference(*args) + assert out.dtype == torch.float8_e4m3fn and out.shape == ref.shape + v_start = (cell.spec["q"] + cell.spec["kv"]) * HEAD_DIM + assert torch.equal(out[:, v_start:].view(torch.uint8), ref[:, v_start:].view(torch.uint8)) + assert _e4m3_steps(out[:, :v_start], ref[:, :v_start]) <= 1 + + +def _e4m3_steps(a, b) -> int: + """The largest distance between two E4M3 tensors in representable values (E4M3 codes order by magnitude).""" + + def signed(t): + c = t.view(torch.uint8).int() + # Sign-magnitude to a line: -0 and +0 both map to 0. + return torch.where(c >= 128, 128 - c, c) + + return int((signed(a) - signed(b)).abs().max()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..4be27dc1a2ff --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_qkv_insert entry: the generic horizontal producer on an FP8 main cache. + +Q, index-Q and every written K / V / index-K cache row are held within one E4M3 step of the fp32 reference. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_qkv_insert import ( # noqa: E501 + m3_qkv_insert as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" + +PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 32, 2, 2, 128, 64 +MAX_POSITIONS = 262144 +EPS = 1e-6 + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +def _args(tokens: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + tokens, (HEADS + 2 * KV + INDEX + 1) * D, device="cuda", generator=g + ).bfloat16() + kv = torch.zeros(PAGES, 2, KV, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + index_k = torch.zeros(PAGES, 1, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + weights = [(torch.randn(D, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + rotary = torch.randn(MAX_POSITIONS, 2, ROTARY // 2, device="cuda", generator=g) + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + return ( + packed, + kv, + None, + index_k, + slots, + None, + HEADS, + KV, + INDEX, + D, + ROTARY, + EPS, + *weights, + rotary, + positions, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _args(cell.spec["m"], zlib.crc32(cell.why.encode())) + certifying(op) + ref = op.reference(*args) + q, index_q = op(*args) + torch.cuda.synchronize() + _within_fp8(q.view(ref["q"].shape), ref["q"], "q") + _within_fp8(index_q.view(ref["index_q"].shape), ref["index_q"], "index_q") + kv, index_k, slots = args[1], args[3], args[4] + for t in range(cell.spec["m"]): + page, within = divmod(int(slots[t]), PAGE) + _within_fp8(kv[page, 0, :, within], ref["k"][t], f"k token {t}") + _within_fp8(kv[page, 1, :, within], ref["v"][t], f"v token {t}") + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..19123d839dea --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_blocks entry of the MiniMax-M3 tp2 target: MSA top-16 block selection with the forced +init / local blocks and per-row valid-block masks.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_blocks import ( # noqa: E501 + m3_select_blocks as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" + +TOPK, INIT, LOCAL = 16, 0, 1 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + n_blocks, total_q = spec["n_blocks"], spec["total_q"] + # Distinct scores (a permutation) so the selection has no ties to break. + scores = ( + torch.randperm(n_blocks * total_q, generator=g, device="cuda") + .float() + .view(1, n_blocks, total_q) + ) + if spec["mixed"]: + n_valid = torch.randint( + 17, n_blocks + 1, (total_q,), generator=g, device="cuda", dtype=torch.int32 + ) + else: + n_valid = torch.full((total_q,), n_blocks, dtype=torch.int32, device="cuda") + with validating(op): + out = op(scores, TOPK, n_valid, INIT, LOCAL, True) + ref = op.reference(scores, TOPK, n_valid, INIT, LOCAL, True) + assert out.shape == ref.shape + assert torch.equal(out, ref) + assert out.permute(1, 0, 2).is_contiguous(), "head-major backing" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..9fc6e361a03e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp2 (Rubin) m3_select_long entry (top-16 MSA block selection for long decode rows, 2 KV heads).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long import ( # noqa: E501 + m3_select_long as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key +# widths, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 2048, 2049, 4096, 4097, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # re-read path). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + # Every row: 10 winners above, ~1500 blocks tied at the 16th score, the rest below; the in-tree op stages up to + # 2048 boundary candidates, so it must break this tie to the smaller block like this op. + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + # Odd tokens: almost everything -inf, so picks run past the finite scores. + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +def _mass_tie_rows(scores, nvb, init_blocks, local_blocks): + """[total_q, num_kv_heads] mask of rows where more than 2048 valid blocks share the 16th-largest score.""" + num_kv_heads, num_blocks, total_q = scores.shape + blocks = torch.arange(num_blocks, device=scores.device) + raw = nvb[:total_q].long() + valid = raw.clamp(0, num_blocks) + forced = scores.permute(2, 0, 1).clone() # [total_q, heads, blocks] + forced = torch.where(blocks < init_blocks, 1.0e30, forced) + forced = torch.where(blocks >= (raw - local_blocks).clamp_min(0)[:, None, None], 1.0e29, forced) + in_row = (blocks[None, :] < valid[:, None])[:, None, :] + ranked = torch.where(in_row, forced, float("-inf")).sort(dim=-1, descending=True).values + kth = ranked[..., min(TOPK, num_blocks) - 1 : min(TOPK, num_blocks)] + return ((forced == kth) & in_row).sum(dim=-1) > 2048 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + + # Selection is an exact claim: the same blocks or a different attention. + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # And it is the in-tree selector's answer, bit for bit (the op the target falls back to) -- except on rows + # where more than 2048 blocks tie at the 16th-largest score: there the in-tree op does not break the tie to + # the smallest blocks, while this op and the reference do. + scores, nvb, _, init_blocks, local_blocks, _ = args + in_tree = select_blocks_from_maxscore( + scores, + topk=TOPK, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + mass_tie = _mass_tie_rows(scores, nvb, init_blocks, local_blocks) + assert torch.equal(out[~mass_tie], in_tree[~mass_tie]), ( + "selected blocks differ from the in-tree minimax_m3_select_blocks" + ) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + """Releasing the dependent grid right after the select's own wait changes only when the next kernel may launch.""" + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + """Releasing the dependent grid once the row's scores are loaded changes only when the next kernel may launch.""" + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) + + +def test_guard_refuses_out_aliasing_scores() -> None: + scores, nvb, _, init_blocks, local_blocks, _ = _case(op.CELLS[0].spec, seed=3) + aliased = ( + scores.view(torch.int32) + .reshape(-1)[: scores.shape[2] * TOPK] + .view(scores.shape[2], 1, TOPK) + ) + with pytest.raises(AssertionError, match="alias"): + with validating(op): + op(scores, nvb, aliased, init_blocks, local_blocks, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py new file mode 100644 index 000000000000..d452cb59f998 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp2 (Rubin) m3_select_long_dsl entry: the long-row top-16 select from one 384-thread CTA per row, +against the reference and bit for bit against the 1024-thread m3_select_long on every cell.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( + m3_select_long_dsl as select_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, + m3_select_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long_dsl import ( # noqa: E501 + m3_select_long_dsl as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key +# widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 1536, 1537, 4224, 4225, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # whole-row re-rank). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # Every row, mass ties included, is the 1024-thread kernel's answer. + scores, nvb, _, init_blocks, local_blocks, _ = args + long_out = torch.empty_like(out) + m3_select_long(scores, nvb, long_out, init_blocks, local_blocks, False) + assert torch.equal(out, long_out), "selected blocks differ from m3_select_long" + + +@pytest.mark.parametrize("layout", ["kv_blocks_q", "kv_q_blocks"]) +@pytest.mark.parametrize("valid", [17, 1294, 2524, 3907, 9216]) +def test_layouts_and_widths_match_the_1024_thread_select(layout, valid) -> None: + """Both score layouts the scorer can hand over, at a valid count in each key width.""" + g = torch.Generator(device="cuda").manual_seed(valid) + nb = max(valid, 129) + s = torch.randn(1, nb, 8, generator=g, device="cuda") + if layout == "kv_q_blocks": + s = s.permute(0, 2, 1).contiguous().permute(0, 2, 1) + nvb = torch.tensor([valid - (t % 2) for t in range(8)], dtype=torch.int32, device="cuda") + out = torch.empty(1, 8, TOPK, dtype=torch.int32, device="cuda").permute(1, 0, 2) + ref = torch.empty_like(out) + op(s, nvb, out, 0, 1, False) + m3_select_long(s, nvb, ref, 0, 1, False) + assert torch.equal(out, ref) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +@pytest.mark.parametrize("early_trigger", [False, True]) +@pytest.mark.parametrize("cell", range(len(op.CELLS))) +def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: + """kernels/m3_select_long_dsl.SELECT_N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the + grid wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and + tied rows).""" + phase = 0 if early_trigger else 1 + spec = op.CELLS[cell].spec + scores, nvb, before, init_blocks, local_blocks, _ = _case(spec, seed=11 + cell) + after = torch.empty_like(before) + select_kernels._compile((True, early_trigger, phase, True))( + scores, nvb, before, init_blocks, local_blocks + ) + select_kernels._compile((True, early_trigger, phase, False))( + scores, nvb, after, init_blocks, local_blocks + ) + assert torch.equal(before, after) + + +def test_graph_replay_selects_the_same() -> None: + spec = op.CELLS[1].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=9) + ref = op.reference(scores, nvb, out, init_blocks, local_blocks) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + for _ in range(3): + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..3fd3f605a230 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,451 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode entry (the target's Triton block-sparse decode and its routes) at the TP2 heads. + +Besides the band of every route around a float64 reference, the Triton merge's MXFP8 o_proj hand-off (published to +``m3_prequant.SLOT``) is held to ``torch.ops.trtllm.mxfp8_quantize(output, True)`` bit for bit: on identity steps and on +selected steps with the long-context range killed (the routes the o_proj consumes when LC2 does not run). +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, # noqa: E501 +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_prequant, + triton_sparse_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + triton_sparse_decode as tsd, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode as stock_decode, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import ( # noqa: E402 — needs the path declared above + HEAD_DIM, + NUM_HEADS, + build_step, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + +# Band, in BF16 ulps (2^-8) of each [32, 128] row's largest magnitude, and of the relative RMS. Every route +# computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference +# is the BF16 rounding on both sides (one ulp of the row at most), the FP16 probabilities before PV, the split-K +# partials kept in FP32 and merged by their log2-sum-exp, and approximate exp2/log2: 2 ulps per element, the two +# roundings (each at most half an ulp) dominating the RMS, so half an ulp. (GB300 TP4: at most 1.80 ulps per element +# and 0.20 ulps RMS over the cells.) +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, output, identity): + return op( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + identity, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=spec["identity"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + certifying(op) + _call(case, output, spec["identity"]) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + spec["identity"], + ) + assert torch.isfinite(output.float()).all(), "padding rows must be zeros, not NaN" + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_flag_does_not_change_the_answer() -> None: + """On an identity table the identity route (the per-row producer) and the selected-table route (the + long-context kernel) compute the same attention.""" + case = build_step((2048, 130), 4, identity=True, seed=17) + outs = [] + for identity in (True, False): + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, identity) + outs.append(output) + assert_within_ulp(outs[0], outs[1], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_guard_refuses_a_padded_selection_slot() -> None: + """The Triton producers index the block table with every one of a token's first min(16, blocks) ids, + -1 included.""" + case = build_step((1500,), 1, identity=True, seed=23) + case.topk[0, 0, 3] = -1 + output = torch.empty(1, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="blocks of the request"): + with validating(op): + _call(case, output, True) + + +def test_guard_refuses_a_false_identity_promise() -> None: + """identity_blocks lets the multi-query producer read one table for the whole request.""" + case = build_step((1500, 1600, 1700, 1800, 1500, 1600, 1700, 1800), 4, identity=True, seed=29) + case.topk[0, 5, :2] = case.topk[0, 5, :2].flip(0) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="identity_blocks promises"): + with validating(op): + _call(case, output, True) + + +def _swizzled_scale_offsets(rows: int, width: int) -> torch.Tensor: + """Offsets of rows [0, rows)'s UE8M0 scales in the 128x4-swizzled layout of mxfp8_quantize(x, True).""" + k_sf = width // 32 + k_tiles = (k_sf + 3) // 4 + r = torch.arange(rows, device="cuda").view(-1, 1) + k = torch.arange(k_sf, device="cuda").view(1, -1) + offsets = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + return offsets.reshape(-1) + + +def _mxfp8_handoff_matches_stock(output) -> None: + """The real rows' E4M3 bytes and UE8M0 scales equal mxfp8_quantize's. The merge writes only the step's rows into + its persistent 128-row scale tile; the padding rows' scales are never read by the o_proj and are not compared.""" + rows = output.shape[0] + flat = output.view(rows, -1) + taken = m3_prequant.SLOT.take(flat) + assert taken is not None, "the route published no MXFP8 hand-off for its output" + q8, sf = taken + stock_q, stock_sf = torch.ops.trtllm.mxfp8_quantize(flat, True) + assert torch.equal( + q8.view(torch.uint8).reshape(rows, -1), stock_q.view(torch.uint8).reshape(rows, -1) + ) + at = _swizzled_scale_offsets(rows, flat.shape[1]) + assert torch.equal( + sf.view(torch.uint8).reshape(-1)[at], stock_sf.view(torch.uint8).reshape(-1)[at] + ) + + +@pytest.mark.parametrize("seq_lens", [(1900,), (2048, 0), (700, 1500)]) +def test_identity_merge_mxfp8_handoff_is_stock_mxfp8_quantize(seq_lens) -> None: + """The Triton merge writes the o_proj hand-off (E4M3 rows, 128x4-swizzled UE8M0 scales) as mxfp8_quantize.""" + case = build_step(seq_lens, 4, identity=True, seed=31 + len(seq_lens)) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(case, output, True) + _mxfp8_handoff_matches_stock(output) + + +@pytest.mark.parametrize("seq_lens", [(40000,), (9000, 40000)]) +def test_killed_long_context_route(seq_lens) -> None: + """With SPARSE_DECODE_LC killed a selected table runs the Triton producer + merge: in the band of the reference, + and its MXFP8 hand-off is mxfp8_quantize of the rows it wrote.""" + case = build_step(seq_lens, 4, identity=False, seed=41 + len(seq_lens)) + landed = triton_sparse_decode.LC_KEYS + triton_sparse_decode.LC_KEYS = frozenset() + try: + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, False) + _mxfp8_handoff_matches_stock(output) + finally: + triton_sparse_decode.LC_KEYS = landed + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_route_is_the_certified_producer_range() -> None: + """Every identity step at the TP2 heads takes the per-row producer, the key modeling.SPARSE_DECODE_PRODUCER + covers (the multi-query producer needs one KV head).""" + assert modeling.SPARSE_DECODE_PRODUCER.covers("per_row") + assert not modeling.SPARSE_DECODE_PRODUCER.covers("multi_query") + for cell in op.CELLS: + keys = tuple(modeling._sparse_decode_producer(cell.spec)) + assert keys == (("per_row",) if cell.spec["identity"] else ()), (cell.why, keys) + + +_IDENTITY_LENS = (1900, 2048, 700, 1500, 1200, 2000, 300, 1024) +_SELECTED_LENS = (9000, 40000, 20000, 5000, 7000, 12000, 30000, 3000) + + +def _stock(case, output): + stock_decode( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + sm_scale=case.sm_scale, + output=output, + decode_query_len=case.decode_query_len, + k_block_scale=case.k_block_scale, + v_block_scale=case.v_block_scale, + k_global_scale=case.k_global_scale, + v_global_scale=case.v_global_scale, + ) + + +@pytest.mark.parametrize("identity", [True, False], ids=["identity", "selected_lc_killed"]) +@pytest.mark.parametrize("requests", [1, 2, 4, 8]) +def test_triton_route_is_bitwise_the_stock_decode(requests, identity) -> None: + """The target's Triton producer + split-K merge writes the stock in-tree MSA decode's rows bit for bit: on SM107 it + takes stock's split-K factor at every row count (4 / 8 / 16 / 32 verify rows here).""" + lens = (_IDENTITY_LENS if identity else _SELECTED_LENS)[:requests] + case = build_step(lens, 4, identity=identity, seed=50 + requests + 100 * identity) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC_KEYS + triton_sparse_decode.LC_KEYS = frozenset() + try: + _call(case, target, identity) + finally: + triton_sparse_decode.LC_KEYS = landed + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)) + + +@pytest.mark.parametrize("requests", [33, 46, 48], ids=["132_rows", "184_rows", "192_rows"]) +def test_plain_route_is_bitwise_the_stock_decode(requests) -> None: + """Above MergeMxfp8.MAX_ROWS (128 rows) a selected table is outside SPARSE_DECODE_LC ("plain") and runs the + Triton producer + merge: the stock decode's rows bit for bit, and no MXFP8 hand-off is published (the o_proj + quantizes its BF16 input itself). 33..48 requests x 4 verify rows: the wide target's steps past one scale tile, + up to WIDE_MAX_TOKENS.""" + case = build_step((3000,) * requests, 4, identity=False, seed=77) + rows = case.q.shape[0] + assert rows > tsd.MergeMxfp8.MAX_ROWS + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _call(case, target, False) + assert m3_prequant.SLOT.take(target.view(rows, -1)) is None + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)) + + +@pytest.mark.parametrize("identity", [True, False], ids=["triton_identity", "lc2_selected"]) +def test_handoff_kill_publishes_nothing_and_keeps_the_rows(identity) -> None: + """MERGE_MXFP8_HANDOFF False: no route publishes the MXFP8 hand-off (the o_proj then quantizes its BF16 input), and + the attention rows are the same bits as with the hand-off.""" + case = build_step((1900,) if identity else (40000,), 4, identity=identity, seed=91 + identity) + on = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + off = torch.empty_like(on) + _call(case, on, identity) + assert m3_prequant.SLOT.take(on.view(on.shape[0], -1)) is not None + triton_sparse_decode.MERGE_MXFP8_HANDOFF = False + try: + _call(case, off, identity) + assert m3_prequant.SLOT.take(off.view(off.shape[0], -1)) is None + finally: + triton_sparse_decode.MERGE_MXFP8_HANDOFF = True + assert torch.equal(on.view(torch.int16), off.view(torch.int16)) + + +def _short_lens(requests: int, g: torch.Generator) -> tuple: + """Short contexts of 1..16 blocks of 128 tokens (the GSM / identity-route regime), a few exactly on a block edge.""" + lens = torch.randint(4, 16 * 128 + 1, (requests,), generator=g).tolist() + if requests > 1: + lens[0] = 128 * int(torch.randint(1, 17, (1,), generator=g)) + return tuple(lens) + + +@pytest.mark.parametrize("seed", range(12)) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_short_table_stress_is_bitwise_the_stock_decode(requests, seed) -> None: + """Short contexts (1..16 blocks) at 1..4 requests x 4 verify rows, the identity route and a random selected route + with -1 tails past each token's blocks: the target's Triton route writes the stock decode's rows bit for bit.""" + g = torch.Generator().manual_seed(1000 * requests + seed) + for identity in (True, False): + case = build_step(_short_lens(requests, g), 4, identity=identity, seed=7 * seed + requests) + target = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC_KEYS + triton_sparse_decode.LC_KEYS = frozenset() + try: + _call(case, target, identity) + finally: + triton_sparse_decode.LC_KEYS = landed + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + identity, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("fault", ["neg_in_window", "id_past_row", "page_out_of_pool"]) +def test_invalid_slot_fails_closed(fault) -> None: + """A corrupt decode input (a -1 inside a token's block window, an id past the block-table row, or a block-table + page outside the pool) is read as an absent block instead of a wild KV read: the rows equal the reference over the + remaining blocks.""" + case = build_step((1500, 2048, 700), 4, identity=False, seed=303) + tok, slot = 5, 2 + if fault == "neg_in_window": + case.topk[1, tok, slot] = -1 + elif fault == "id_past_row": + case.topk[1, tok, slot] = case.block_table.shape[1] + 7 + else: + req = tok // 4 + blk = int(case.topk[1, tok, slot]) + bad_page = case.k_paged.shape[0] + 1000 + case.block_table = case.block_table.clone() + case.block_table[req, blk] = bad_page + # Every token of the request that selected that block loses it. + hits = (case.topk[:, req * 4 : req * 4 + 4] == blk).nonzero().tolist() + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + landed = triton_sparse_decode.LC_KEYS + triton_sparse_decode.LC_KEYS = frozenset() + try: + _call(case, output, False) + finally: + triton_sparse_decode.LC_KEYS = landed + torch.cuda.synchronize() + # The reference walks each token's first min(16, blocks) entries and drops -1: the corrupt slots as -1. + ref_topk = case.topk.clone() + if fault == "page_out_of_pool": + for h, t, k in hits: + ref_topk[h, req * 4 + t, k] = -1 + else: + ref_topk[1, tok, slot] = -1 + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + ref_topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert torch.isfinite(output.float()).all() + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def _identity_table( + seq_lens: torch.Tensor, decode_query_len: int, heads: int, lag: int = 0 +) -> torch.Tensor: + """The identity selection of every verify token, staged from lengths `lag` tokens behind seq_lens.""" + seq = (seq_lens.cpu().to(torch.int64) - lag).clamp(min=0).repeat_interleave(decode_query_len) + intra = torch.arange(seq.shape[0]) % decode_query_len + n_valid = ((seq - decode_query_len + intra + 1).clamp(min=0) + 127) // 128 + ar = torch.arange(16) + table = torch.where(ar[None, :] < n_valid[:, None], ar[None, :], -1).to(torch.int32) + return table.unsqueeze(0).expand(heads, -1, -1).contiguous().cuda() + + +# Each crosses a 128-token edge within its four verify tokens or between the stale and the true lengths. +_EDGE_LENS = (1284, 1287, 1281, 129, 132, 2048, 1028, 5, 643, 1920, 256, 1) + + +@pytest.mark.parametrize("fault", ["stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_identity_route_reads_its_own_lengths_not_the_table(requests, fault) -> None: + """The identity route derives block i for slot i from the kv lengths the kernel reads: a staged table that lags + the lengths by a verify step (a 128-token edge crossed), or is garbage, gives the stock decode's rows over the + true identity selection bit for bit.""" + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=400 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + assert torch.equal(case.topk, true_table), ( + "build_step's identity table is the identity selection" + ) + if fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _stock(case, stock) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + _call(case_staged, target, True) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py new file mode 100644 index 000000000000..b726d37e1dd3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for trtllm::minimax_m3_sm107_tp2_sparse_decode_lc (one-launch long-context sparse decode, TP2). + +Two contracts: +- the BF16 rows sit in the same band around a float64 reference as the stock Triton sparse decode, and every variant + (linear softmax, pre-wait block-table staging, MXFP8 hand-off on or off) writes the same rows as the stock + log-domain recurrence where the arithmetic is shared; +- the MXFP8 hand-off (q8, sf) is bitwise torch.ops.trtllm.mxfp8_quantize(out, True) of the rows the kernel wrote: the + E4M3 bytes row-major and the UE8M0 scales in its 128x4-swizzled layout. The o_proj consumes this as its + prequantized activation, so this test is the interface contract. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + m3_sparse_decode_lc as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + mxfp8_handoff_reference, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402, F401 — needs the path declared above + assert_within_ulp, + certifying, +) + +PAGE, DIM, TOPK, GROUP, KV_HEADS = 128, 128, 16, 16, 2 +WIDTH = KV_HEADS * GROUP * DIM +SM_SCALE = DIM**-0.5 +# Band, in BF16 ulps (2^-8) of each output row's largest magnitude, and of the relative RMS. Same as the sm_103 TP4 +# kernel; on SM107 the stock Triton decode and this kernel both measure at most 1.03 ulps and 0.26 RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + +_E2M1 = [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6] + + +def _op(): + return torch.ops.trtllm.minimax_m3_sm107_tp2_sparse_decode_lc + + +def _e4m3_bytes(shape, gen, lo=0x28, hi=0x48): + mag = torch.randint(lo, hi, shape, dtype=torch.int32, device="cuda", generator=gen) + sign = torch.randint(0, 2, shape, dtype=torch.int32, device="cuda", generator=gen) << 7 + return (mag | sign).to(torch.uint8) + + +def _case(seq_lens, seed, dql=4, overlap=14, q_hi=0x48, tie=False): + """A packed-NVFP4 paged cache, ascending top-16 tables with a -1 tail, and FP8 queries for TP2 geometry.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + reqs, total_q = len(seq_lens), len(seq_lens) * dql + nblk = [(s + PAGE - 1) // PAGE for s in seq_lens] + pages = sum(nblk) + 8 + k = torch.randint( + 0, 256, (pages, KV_HEADS, PAGE, DIM // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + v = torch.randint( + 0, 256, (pages, KV_HEADS, PAGE, DIM // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + if tie: + k[:] = k[:1, :, :1, :] + ksc = _e4m3_bytes((pages, KV_HEADS, PAGE, DIM // 16), gen) + vsc = _e4m3_bytes((pages, KV_HEADS, PAGE, DIM // 16), gen) + perm = torch.randperm(pages, device="cuda", generator=gen).int() + bt = torch.full((reqs, max(nblk)), -1, dtype=torch.int32, device="cuda") + start = 0 + for r, n in enumerate(nblk): + bt[r, :n] = perm[start : start + n] + start += n + topk = torch.full((KV_HEADS, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for r, s in enumerate(seq_lens): + for h in range(KV_HEADS): + last = (s + PAGE - 1) // PAGE + base = torch.randperm(max(last - 1, 1), device="cuda", generator=gen)[ + : TOPK - 1 + ].tolist() + for qo in range(dql): + nb = (s - dql + qo + 1 + PAGE - 1) // PAGE + sel = [b for b in base[:overlap] if b < nb - 1] + for b in torch.randperm(max(nb - 1, 1), device="cuda", generator=gen).tolist(): + if len(sel) >= min(TOPK, nb) - 1: + break + if b not in sel and b < nb - 1: + sel.append(b) + sel = sorted(set(sel + [nb - 1]))[: min(TOPK, nb)] + topk[h, r * dql + qo, : len(sel)] = torch.tensor( + sel, dtype=torch.int32, device="cuda" + ) + q = _e4m3_bytes((total_q, KV_HEADS * GROUP, DIM), gen, lo=0x20, hi=q_hi).view( + torch.float8_e4m3fn + ) + return dict( + q=q, + k=k, + v=v, + ksc=ksc, + vsc=vsc, + kgs=torch.tensor([0.03], device="cuda"), + vgs=torch.tensor([0.05], device="cuda"), + topk=topk, + bt=bt, + seq=torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + dql=dql, + ) + + +def _dequant(packed, sc, quad): + lut = torch.tensor(_E2M1, dtype=torch.float64, device="cuda") + vals = torch.stack([lut[(packed & 0xF).long()], lut[(packed >> 4).long()]], dim=-1).reshape( + *packed.shape[:3], DIM + ) + if ( + quad + ): # V scales: token t, column c at byte (t // 4) * 32 + t % 4 + 4 * c within the page/head + t = torch.arange(PAGE, device="cuda")[:, None] + idx = (t // 4) * 32 + t % 4 + 4 * torch.arange(8, device="cuda")[None, :] + sc = sc.reshape(*sc.shape[:2], PAGE * 8)[:, :, idx] + return vals * sc.contiguous().view(torch.float8_e4m3fn).to(torch.float64).repeat_interleave( + 16, dim=-1 + ) + + +def _reference(c): + q = c["q"].to(torch.float64) + out = torch.zeros(q.shape, dtype=torch.float64, device="cuda") + dql = c["dql"] + for tok in range(q.shape[0]): + r, qo = divmod(tok, dql) + kv_len = max(int(c["seq"][r]) - dql + qo + 1, 0) + for h in range(KV_HEADS): + blks = [b for b in c["topk"][h, tok].tolist() if b >= 0][ + : min(TOPK, (kv_len + PAGE - 1) // PAGE) + ] + if not blks: + continue + pg = c["bt"][r, blks].long() + kk = _dequant(c["k"][pg, h : h + 1], c["ksc"][pg, h : h + 1], False)[:, 0].reshape( + -1, DIM + ) + vv = _dequant(c["v"][pg, h : h + 1], c["vsc"][pg, h : h + 1], True)[:, 0].reshape( + -1, DIM + ) + pos = ( + torch.tensor(blks, device="cuda")[:, None] * PAGE + + torch.arange(PAGE, device="cuda") + ).reshape(-1) + s = (q[tok, h * GROUP : (h + 1) * GROUP] @ kk.T) * (SM_SCALE * c["kgs"].item()) + s[:, pos >= kv_len] = float("-inf") + out[tok, h * GROUP : (h + 1) * GROUP] = (torch.softmax(s, dim=-1) @ vv) * c[ + "vgs" + ].item() + return out.reshape(q.shape[0], WIDTH) + + +def _band(x, ref): + x = x.to(torch.float64).reshape(-1, DIM) + ref = ref.reshape(-1, DIM) + rowmax = ref.abs().amax(dim=-1, keepdim=True) + live = rowmax[:, 0] > 0 + unit = (rowmax * 2.0**-8).clamp_min(1e-300) + err = ((x - ref).abs() / unit)[live] + rms = (((x - ref) ** 2).mean(dim=-1).sqrt() / unit[:, 0])[live] + return (float(err.max()) if err.numel() else 0.0), (float(rms.max()) if rms.numel() else 0.0) + + +def _run(c, linear=False, prewait=False, mxfp8=False, n1=None, use_pdl=False, out=None): + total_q = c["q"].shape[0] + if out is None: + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + q8 = torch.zeros(total_q, WIDTH, dtype=torch.uint8, device="cuda") if mxfp8 else None + sf_bytes = 128 * (WIDTH // 32) * ((total_q + 127) // 128) + sf = torch.zeros(sf_bytes, dtype=torch.uint8, device="cuda") if mxfp8 else None + _op()( + c["q"].view(torch.uint8), + c["k"], + c["v"], + c["ksc"], + c["vsc"], + c["kgs"], + c["vgs"], + c["topk"], + c["bt"], + c["seq"], + out, + q8, + sf, + WIDTH // 128, + c["dql"], + SM_SCALE, + linear, + use_pdl, + prewait, + n1, + ) + return out, q8, sf + + +_CASES = { + "b1-165k": dict(seq_lens=[165521]), + "b2-mixed": dict(seq_lens=[165521, 98765]), + "b2-short": dict(seq_lens=[130, 5000]), + "b1-one-token-tail": dict(seq_lens=[128 * 40 + 1]), + "b1-padded-topk": dict(seq_lens=[700]), + "b1-tiny": dict(seq_lens=[4]), + "b2-large-logits": dict(seq_lens=[165521, 64000], q_hi=0x70), + "b1-ties": dict(seq_lens=[20000], tie=True), + "b2-disjoint": dict(seq_lens=[165521, 165521], overlap=0), + "b4": dict(seq_lens=[165521, 120000, 80000, 30000]), +} + + +@pytest.mark.parametrize("name", list(_CASES)) +@pytest.mark.parametrize("seed", [0, 1]) +def test_band_vs_float64(name, seed) -> None: + c = _case(seed=seed, **_CASES[name]) + out, _, _ = _run(c) + elem, rms = _band(out, _reference(c)) + assert torch.isfinite(out.float()).all() + assert elem <= _ELEMENT_ULP and rms <= _RMS_ULP, (elem, rms) + + +@pytest.mark.parametrize("name", ["b1-165k", "b2-mixed", "b1-padded-topk", "b4"]) +def test_variants_write_the_same_rows(name) -> None: + c = _case(seed=7, **_CASES[name]) + base, _, _ = _run(c) + for prewait in (False, True): + for mxfp8 in (False, True): + out, _, _ = _run(c, prewait=prewait, mxfp8=mxfp8) + assert torch.equal(out.view(torch.int16), base.view(torch.int16)), (prewait, mxfp8) + lin, _, _ = _run(c, linear=True) + elem, rms = _band(lin, _reference(c)) + assert elem <= _ELEMENT_ULP and rms <= _RMS_ULP, (elem, rms) + + +@pytest.mark.parametrize("name", list(_CASES)) +@pytest.mark.parametrize("linear", [False, True]) +def test_mxfp8_handoff_is_stock_mxfp8_quantize(name, linear) -> None: + """The o_proj interface: (q8, sf) == mxfp8_quantize(out, is_sf_swizzled_layout=True), bitwise.""" + c = _case(seed=3, **_CASES[name]) + out, q8, sf = _run(c, linear=linear, mxfp8=True) + ref_q, ref_sf = torch.ops.trtllm.mxfp8_quantize(out, True) + assert torch.equal(q8, ref_q.view(torch.uint8).reshape(q8.shape)) + ref_sf = ref_sf.view(torch.uint8).reshape(-1) + assert torch.equal(sf[: ref_sf.numel()], ref_sf) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The catalog entry's cells through its own __call__ (guarded by is_valid) and its own reference.""" + spec = cell.spec + c = _case( + list(spec["seq_lens"]), + seed=zlib.crc32(cell.why.encode()) % 2**31, + dql=spec["decode_query_len"], + ) + total_q = c["q"].shape[0] + # the hand-off buffers hold whole 128-row tiles (MergeMxfp8.MAX_ROWS) + alloc = (total_q + 127) // 128 * 128 + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + q8 = torch.zeros(alloc, WIDTH, dtype=torch.uint8, device="cuda") if spec["mxfp8"] else None + sf = ( + torch.zeros(alloc * (WIDTH // 32), dtype=torch.uint8, device="cuda") + if spec["mxfp8"] + else None + ) + args = ( + c["q"].view(torch.uint8), + c["k"], + c["v"], + c["ksc"], + c["vsc"], + c["kgs"], + c["vgs"], + c["topk"], + c["bt"], + c["seq"], + out, + q8, + sf, + WIDTH // 128, + c["dql"], + SM_SCALE, + spec["linear_softmax"], + False, + spec["prewait"], + ) + certifying(op) + op(*args) + ref = op.reference(*args) + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + assert torch.isfinite(out.float()).all() + if spec["mxfp8"]: + ref_q8, ref_sf = mxfp8_handoff_reference(out, alloc) + assert torch.equal(q8[:total_q], ref_q8) + assert torch.equal(sf, ref_sf) + stock_q, stock_sf = torch.ops.trtllm.mxfp8_quantize(out, True) + assert torch.equal(q8[:total_q], stock_q.view(torch.uint8).reshape(total_q, WIDTH)) + assert torch.equal(sf[: stock_sf.numel()], stock_sf.view(torch.uint8).reshape(-1)) + + +# LC on vs the killed range (SPARSE_DECODE_LC = frozenset(): the stock Triton producer + merge) through the target's +# minimax_m3_sparse_attn_decode on the same inputs. Not bitwise: the split-K merge sums the 16 partials in Triton's +# reduction order there and in chunk order here (the per-block math is the same), so the rows differ by summation +# order only. Band, in BF16 ulps of each row's largest magnitude: measured on SM107 at most 1.92 per element and +# 0.18 RMS (R4/R13 of the MSA lane), so 2 / 0.25. +_LC_VS_KILLED_ELEMENT_ULP = 2.0 +_LC_VS_KILLED_RMS_ULP = 0.25 + + +def _target_decode(c, killed): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_prequant, + triton_sparse_decode, + ) + + landed = triton_sparse_decode.LC_KEYS + if killed: + triton_sparse_decode.LC_KEYS = frozenset() + try: + total_q = c["q"].shape[0] + out = torch.empty(total_q, NUM_HEADS_TOTAL, DIM, dtype=torch.bfloat16, device="cuda") + triton_sparse_decode.minimax_m3_sparse_attn_decode( + c["q"], + c["k"], + c["v"], + c["topk"], + c["bt"], + c["seq"], + sm_scale=SM_SCALE, + output=out, + decode_query_len=c["dql"], + k_block_scale=c["ksc"], + v_block_scale=c["vsc"], + k_global_scale=c["kgs"], + v_global_scale=c["vgs"], + ) + torch.cuda.synchronize() + finally: + triton_sparse_decode.LC_KEYS = landed + m3_prequant.SLOT.clear() + return out.view(total_q, WIDTH) + + +NUM_HEADS_TOTAL = KV_HEADS * GROUP + + +@pytest.mark.parametrize("name", list(_CASES)) +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_lc_vs_killed_stock_path(name, seed) -> None: + """The landed cell and its kill switch on the same step: rows within summation-order distance of each other.""" + c = _case(seed=seed, **_CASES[name]) + lc = _target_decode(c, killed=False) + stock = _target_decode(c, killed=True) + assert torch.isfinite(lc.float()).all() and torch.isfinite(stock.float()).all() + elem, rms = _band(lc, stock.to(torch.float64)) + assert elem <= _LC_VS_KILLED_ELEMENT_ULP and rms <= _LC_VS_KILLED_RMS_ULP, (elem, rms) + ref = _reference(c) + for out in (lc, stock): + e, r = _band(out, ref) + assert e <= _ELEMENT_ULP and r <= _RMS_ULP, (e, r) + + +# N1 (n1_ids): the layer's previous-step block ids, prefetched into L2 before the grid wait and overwritten with this +# step's. A hint: the rows must be bitwise the rows without it, whatever the buffer holds. + + +def _n1_expected(c, rows, before=None): + """The buffer after a launch: row ``tok`` < T holds topk[h, tok, r] for r below the token's block count, else -1; + rows at or past T keep what they held (``before``, default -1).""" + total_q, dql = c["q"].shape[0], c["dql"] + exp = ( + before.clone() + if before is not None + else torch.full((KV_HEADS, rows, TOPK), -1, dtype=torch.int32, device="cuda") + ) + for tok in range(min(total_q, rows)): + r, qo = divmod(tok, dql) + kv_len = max(int(c["seq"][r]) - dql + qo + 1, 0) + real = min(TOPK, (kv_len + PAGE - 1) // PAGE) + exp[:, tok] = -1 + exp[:, tok, :real] = c["topk"][:, tok, :real] + return exp + + +def _n1_garbage(rows, seed): + gen = torch.Generator(device="cuda").manual_seed(seed) + ids = torch.randint( + -5, 1 << 20, (KV_HEADS, rows, TOPK), dtype=torch.int32, device="cuda", generator=gen + ) + ids[0, :, 0] = 2**31 - 1 + ids[1, :, 0] = -(2**31) + return ids + + +# In-model request counts (1, 2, 4, 5, 8 Eagle3 verify requests of 4 rows), a full 128-row step, a plain decode. +_N1_CASES = { + "m1": dict(seq_lens=[165521]), + "m2": dict(seq_lens=[165521, 98765]), + "m4": dict(seq_lens=[165521, 120000, 700, 30000]), + "m5": dict(seq_lens=[165521, 4, 130, 98765, 5000]), + "m8": dict(seq_lens=[165521, 98765, 64000, 30000, 20000, 5000, 700, 130]), + "m32": dict(seq_lens=[4096 * (i + 1) + 77 * i for i in range(32)]), + "m2-dql1": dict(seq_lens=[165521, 5000], dql=1), +} + + +@pytest.mark.parametrize("name", list(_N1_CASES)) +def test_n1_rows_bitwise_and_ids_recorded(name) -> None: + c = _case(seed=11, **_N1_CASES[name]) + total_q = c["q"].shape[0] + for prewait in (False, True): + for mxfp8 in (False, True): + base, bq8, bsf = _run(c, prewait=prewait, mxfp8=mxfp8) + ids = torch.full((KV_HEADS, 128, TOPK), -1, dtype=torch.int32, device="cuda") + # First step (all -1: no prefetch), second (last step's ids: prefetch), then stale garbage ids. + for step, init in enumerate((None, None, _n1_garbage(128, 3))): + if init is not None: + ids.copy_(init) + exp = _n1_expected(c, 128, ids) + out, q8, sf = _run(c, prewait=prewait, mxfp8=mxfp8, n1=ids) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), base.view(torch.int16)), ( + prewait, + mxfp8, + step, + ) + if mxfp8: + assert torch.equal(q8, bq8) and torch.equal(sf, bsf), (prewait, mxfp8, step) + assert torch.equal(ids, exp), (prewait, mxfp8, step) + if total_q > 1: + # A buffer with fewer rows than the step: those rows record, the rest skip both prefetch and record. + rows = total_q // 2 + ids = torch.full((KV_HEADS, rows, TOPK), -1, dtype=torch.int32, device="cuda") + out, _, _ = _run(c, n1=ids) + assert torch.equal(out.view(torch.int16), _run(c)[0].view(torch.int16)) + assert torch.equal(ids, _n1_expected(c, rows)) + + +def test_n1_pdl_stress_and_graph_replay() -> None: + """Back-to-back PDL launches of two layers (each with its own id buffer, alternating steps), eager and replayed + from a CUDA graph: every launch writes the rows of the launch without N1, and the buffers end at the last step's + ids.""" + cases = [_case(seed=21, **_N1_CASES["m2"]), _case(seed=22, **_N1_CASES["m5"])] + bases = [_run(c, mxfp8=True)[0] for c in cases] + ids = [torch.full((KV_HEADS, 128, TOPK), -1, dtype=torch.int32, device="cuda") for _ in cases] + iters = 200 + outs = [[torch.empty_like(b) for _ in range(iters)] for b in bases] + for i in range(iters): + for j, c in enumerate(cases): + _run(c, mxfp8=True, n1=ids[j], use_pdl=True, out=outs[j][i]) + torch.cuda.synchronize() + for j in range(len(cases)): + assert all(torch.equal(o.view(torch.int16), bases[j].view(torch.int16)) for o in outs[j]), j + assert torch.equal(ids[j], _n1_expected(cases[j], 128)), j + for t in ids: + t.fill_(-1) + gouts = [torch.empty_like(b) for b in bases] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream): + for j, c in enumerate(cases): # warm the launch configuration outside capture + _run(c, mxfp8=True, n1=ids[j], use_pdl=True, out=gouts[j]) + torch.cuda.synchronize() + with torch.cuda.graph(graph, stream=stream): + for _ in range(4): + for j, c in enumerate(cases): + _run(c, mxfp8=True, n1=ids[j], use_pdl=True, out=gouts[j]) + garbage = _n1_garbage(128, 5) + for t in ids: + t.copy_(garbage) + for _ in range(25): + for o in gouts: + o.zero_() + graph.replay() + torch.cuda.synchronize() + for j in range(len(cases)): + assert torch.equal(gouts[j].view(torch.int16), bases[j].view(torch.int16)), j + assert torch.equal(ids[j], _n1_expected(cases[j], 128, garbage)), j diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py new file mode 100644 index 000000000000..6d4616218e0b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for kernels/m3_tokmap_guard: a one-request multi-token call of +thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a +stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the +appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for +every length. Without the guard the stale map moves rows (the negative control), and a several-request call passes +through with its own map untouched. + +The fresh-map reference forces the op's own build (a scaled RoPE type with no inv-freq cache, at rotary dim 0, so no +row is rotated either way). +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_tokmap_guard as guard, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = 32, 2, 128, 128 # the TP2 drafter shard (64 / 2 heads, 4 / 2 KV heads) +BUCKET = 4 +_PEER_EXTRA = 777 # the second request's extra length behind the first +FP8_KV_QUANT_MODE = 128 +_DEV = torch.device("cuda") # unindexed on purpose: install() and the lookup normalize it +_TAKEN = [] + + +@pytest.fixture(scope="module", autouse=True) +def _installed(): + guard.install(_DEV, 48, _TAKEN.append) + yield + + +class _Step: + """One request (or ``nreq``) of BUCKET padded rows over a fresh FP8 paged cache, persistent across calls like a + real step's buffers.""" + + def __init__(self, kv_len: int, nreq: int = 1, seed: int = 0): + g = torch.Generator(device=_DEV).manual_seed(seed + kv_len) + self.nreq = nreq + # Per-request inputs (page tables, generation lengths, position offsets) for at least two requests, the page + # tables at the peer's longer length (_seq_with_peer): the stale map sends rows to request 1, which must have + # real entries behind it, also on the unguarded calls. + tables = max(nreq, 2) + pages = math.ceil((kv_len + _PEER_EXTRA + BUCKET) / PAGE) + self.nblocks = tables * pages + 2 + self.pool0 = ( + torch.randn(self.nblocks, 2, KVH, PAGE, D, device=_DEV, generator=g) + .clamp_(-4, 4) + .to(torch.float8_e4m3fn) + ) + self.pool = self.pool0.clone() + perm = torch.randperm(self.nblocks, device=_DEV, generator=g)[: tables * pages].view( + tables, pages + ) + self.offsets = torch.zeros(1, tables, 2, pages, dtype=torch.int32, device=_DEV) + self.offsets[0, :, 0] = 2 * perm + self.offsets[0, :, 1] = 2 * perm + 1 + self.pointers = torch.tensor([[self.pool.data_ptr(), 0]], dtype=torch.int64) + self.mapping = torch.zeros(1, 2, dtype=torch.int32) + self.qkv = torch.randn( + nreq * BUCKET, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=_DEV, generator=g + ) + # Request lengths after this step's append; generation lengths (accepted rows) set per case. + self.seq = torch.full((nreq,), kv_len, dtype=torch.int32, device=_DEV) + self.gen = torch.full((tables,), BUCKET, dtype=torch.int32, device=_DEV) + self.pos = ( + torch.arange(BUCKET, dtype=torch.int32, device=_DEV).repeat(tables, 1).contiguous() + ) + self.one = torch.ones(1, dtype=torch.float32, device=_DEV) + self.ws = torch.empty(64 << 20, dtype=torch.uint8, device=_DEV) + self.kv_len = kv_len + + def call(self, op, fresh_map: bool = False): + """One generation preprocess; returns a copy of the processed query.""" + out = op( + self.qkv, + self.ws, + self.seq, + self.gen, + self.pos, + self.offsets, + self.pointers, + self.mapping, + self.one, + self.one, + None, + None, + None, + None, + 0, + 0, + H, + KVH, + D, + PAGE, + FP8_KV_QUANT_MODE, + 1 << 20, + 1 << 20, + self.nreq * BUCKET, + self.nreq, + BUCKET, + self.kv_len, + 0, + 10000.0, + 1 if fresh_map else 0, + 1.0, + 1 << 20, + 0, + 1.0, + 1.0, + False, + BUCKET, + 0, + torch.cuda.get_device_properties(0).multi_processor_count, + self.nblocks, + 2, + True, + False, + ) + return out[0].clone() + + def poison_map(self) -> None: + """A stale multi-request map in the op's one-request map slice: rows 1.. on request 1.""" + layout = thop.get_trtllm_gen_generation_workspace_layout( + self.qkv.dtype, 1, self.nreq * BUCKET, H, D, 0, KVH, 0, False, False + ) + off = int(layout["tokens_info_offset"]) + stale = torch.tensor( + [[0, 0]] + [[1, t] for t in range(BUCKET - 1)], dtype=torch.int32, device=_DEV + ) + self.ws[off : off + 8 * BUCKET].view(torch.int32).view(BUCKET, 2).copy_(stale) + + def reset(self) -> None: + self.pool.copy_(self.pool0) + + +def _seq_with_peer(step: _Step) -> None: + # A second request's length behind the first, so a row the stale map moves to request 1 reads a real length. + step.seq = torch.tensor( + [step.kv_len, step.kv_len + _PEER_EXTRA], dtype=torch.int32, device=_DEV + ) + + +@pytest.mark.parametrize("kv_len", [130, 8191, 34117]) +@pytest.mark.parametrize("n", [1, 2, 3, 4]) +def test_guard_equals_a_fresh_map(kv_len, n) -> None: + step = _Step(kv_len) + _seq_with_peer(step) + step.gen.fill_(n) + step.poison_map() + step.reset() + q_ref = step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + + step.poison_map() + step.reset() + q = step.call(thop.trtllm_gen_generation_preprocess) + # The processed query's n live rows (rows past the request's length are padding the op leaves unwritten). + assert torch.equal(q[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)) + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + assert BUCKET in _TAKEN + + +def test_stale_map_moves_rows_without_the_guard() -> None: + step = _Step(34117) + _seq_with_peer(step) + step.gen.fill_(BUCKET) + step.poison_map() + step.reset() + step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + step.poison_map() + step.reset() + step.call(guard._original) + assert not torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + + +def test_one_captured_graph_serves_every_length() -> None: + step = _Step(34000) + _seq_with_peer(step) + refs = {} + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + q = step.call(guard._original, fresh_map=True) + refs[n] = (q, step.pool.clone()) + step.gen.fill_(BUCKET) + step.call(thop.trtllm_gen_generation_preprocess) # eager first, as the warmup runs it + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = thop.trtllm_gen_generation_preprocess + q_static = step.call(out) + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + g.replay() + torch.cuda.synchronize() + q_ref, pool_ref = refs[n] + assert torch.equal(q_static[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)), n + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)), n + + +def test_several_requests_pass_through() -> None: + step = _Step(8191, nreq=2) + layout = thop.get_trtllm_gen_generation_workspace_layout( + step.qkv.dtype, 2, 2 * BUCKET, H, D, 0, KVH, 0, False, False + ) + step.reset() + step.call(thop.trtllm_gen_generation_preprocess) + pool_guarded = step.pool.clone() + step.reset() + step.call(guard._original) + assert torch.equal(step.pool.view(torch.uint8), pool_guarded.view(torch.uint8)) + assert int(layout["tokens_info_size"]) >= 8 * 2 * BUCKET diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..cf8c5ab15574 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the msa_kv_update entry (the one-launch MSA KV-length correction).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.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_107__tp2.ops.attention.msa_kv_update import ( # noqa: E501 + TOKENS_PER_BLOCK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" + +#: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. +_SPARE = 3 + + +def _case(spec, seed): + """The metadata's buffers for one step, with stale values in every output so a skipped write shows.""" + g = torch.Generator(device="cpu").manual_seed(seed) + batch, qlen = spec["batch"], spec["query_len"] + total_q = batch * qlen + staged = torch.randint(qlen, 5000, (batch,), generator=g, dtype=torch.int32) + accept = spec["accept"] + if accept == "full": + kv = staged.clone() + elif accept == "page_edge": + staged = torch.tensor([130, 257, 2 * TOKENS_PER_BLOCK + qlen][:batch], dtype=torch.int32) + kv = staged - torch.tensor([0, 1, qlen - 1][:batch], dtype=torch.int32) + else: + kv = staged - torch.randint(0, qlen, (batch,), generator=g, dtype=torch.int32) + if accept == "padded": + kv[batch // 2 :] = 0 + staged[batch // 2 :] = 0 + ncols = (5000 + qlen) // TOKENS_PER_BLOCK + 2 + table = ( + torch.randperm(batch * ncols + 7, generator=g)[: batch * ncols] + .to(torch.int32) + .reshape(batch, ncols) + ) + dev = dict(device="cuda", dtype=torch.int32) + stale = -12345 + return ( + torch.cat([kv, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.cat([staged, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.full((batch + _SPARE,), qlen, **dev), + torch.arange(batch, **dev).repeat_interleave(qlen), + torch.arange(qlen, **dev).repeat(batch), + table.cuda(), + TOKENS_PER_BLOCK, + batch, + total_q, + torch.full((batch + _SPARE,), stale, **dev), + torch.full((batch + 1 + _SPARE,), stale, **dev) if spec["cu_kv"] else None, + torch.full((total_q + _SPARE,), stale, **dev), + torch.full((total_q + _SPARE,), stale, **dev) if spec["n_valid"] else None, + ) + + +def _check(args, ref): + seq, cu_kv, loc, n_valid = args[9], args[10], args[11], args[12] + batch, total_q = args[7], args[8] + ref_seq, ref_cu, ref_loc, ref_nv = ref + # Integer bookkeeping: exact, and nothing past the live rows is touched. + assert torch.equal(seq[:batch], ref_seq) + assert torch.equal(loc[:total_q], ref_loc) + assert bool((seq[batch:] == -12345).all()) and bool((loc[total_q:] == -12345).all()) + assert (cu_kv is None) == (ref_cu is None) and (n_valid is None) == (ref_nv is None) + if cu_kv is not None: + assert torch.equal(cu_kv[: batch + 1], ref_cu) + assert bool((cu_kv[batch + 1 :] == -12345).all()) + if n_valid is not None: + assert torch.equal(n_valid[:total_q], ref_nv) + assert bool((n_valid[total_q:] == -12345).all()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + op(*args) + _check(args, op.reference(*args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_hook_takes_every_cell(cell) -> None: + """The target's hook (what the metadata calls) takes each certified step and writes what the op does.""" + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + 1) + assert hook(*args) is True + _check(args, op.reference(*args)) + + +def test_the_hook_declines_what_the_op_would_misread() -> None: + """A strided operand is read wrongly by the op, so the hook declines it (the metadata then runs its own + torch update) and writes nothing.""" + args = list(_case(op.CELLS[2].spec, seed=5)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + before = args[11].clone() + assert hook(*args) is False + assert torch.equal(args[11], before) + + +def test_guard_refuses_a_strided_operand() -> None: + args = list(_case(op.CELLS[2].spec, seed=6)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_index_past_the_batch() -> None: + args = list(_case(op.CELLS[2].spec, seed=7)) + args[3] = args[3].clone() + args[3][-1] = args[7] + with pytest.raises(AssertionError, match="q_batch_row"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py new file mode 100644 index 000000000000..822a3b2a0ffc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the ``m3_allreduce`` entry, 2 ranks (SM107): its cells against the entry's reference, +and the fused draft loop's layer boundaries built on it (``M3Eagle3Draft.ar_norm``: the plain all-reduce, then +FlashInfer's fused add + RMSNorm) against stock Eagle3's, bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy the single-node drafter can run (AUTO, +ONESHOT, TWOSHOT, NCCL; MNNVL is multi-node only), at tp2 shapes (H = 6144, eps = 1e-6, 1 to 32 rows). The inputs are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own 2-rank launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_allreduce.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py +""" + +import math +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 4, 8, 32) +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +RANK = 0 +WORLD = 1 + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on each rank. + parts = [x / 2, x / 2] + return ( + [p.bfloat16() for p in parts[:WORLD]], + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A + B, -A - B + d) onto a unit residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a + b, (-a - b + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return [p.bfloat16() for p in parts[:WORLD]], residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +def _ulps_at(out: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the magnitude ``scale`` (the operands'), not at the result's.""" + return ((out.double() - ref.double()).abs() / _ulp_at(scale.abs())).max().item() + + +def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary (``M3Eagle3Draft.ar_norm``): ``m3_allreduce`` plain, then the catalog's + ``flashinfer_fused_add_rmsnorm``.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + + return flashinfer_fused_add_rmsnorm(m3_allreduce(x, ar, None), residual.clone(), w, EPS) + + +def _check_cells(mapping: Any) -> None: + """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum + to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude + u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final + roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). + The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm.functional import AllReduceStrategy + + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating + + ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + for i, cell in enumerate(m3_allreduce.CELLS): + spec = cell.spec + xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) + params = ( + None + if spec.get("plain") + else AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual.clone(), + norm_weight=w, + eps=EPS, + trigger_completion_at_end=False, + ) + ) + # The reference first: the fused path may write the updated residual into ``params.residual``. + ref = m3_allreduce.reference(xs, ar, params) + with validating(m3_allreduce): + out = m3_allreduce(xs[RANK], ar, params) + torch.cuda.synchronize() + total = sum(x.double() for x in xs) + if params is None: + err = _ulps_at(out, ref, total) + assert err <= 1.0, (cell.why, err) + else: + err = _ulps_at(out[1], ref[1], torch.maximum(total.abs(), residual.double().abs())) + assert err <= 2.5, (cell.why, "residual", err) + h = out[1].double() + own = h * torch.rsqrt(h.square().mean(-1, keepdim=True) + EPS) * w.double() + assert _ulps(out[0], own) <= 2.0, (cell.why, "normed", _ulps(out[0], own)) + dist.barrier() + if RANK == 0: + print(f"{RESULT_TAG} cell {cell.why} passed", flush=True) + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + _check_cells(mapping) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + ulp, res_ulp = 0.0, 0.0 + for m in TOKENS: + xs, residual, w = build(m) + target_out, target_res = _target(ar, xs[RANK], residual, w) + stock_out, stock_res = _stock(ar, xs[RANK], residual, w) + torch.cuda.synchronize() + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {name} m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {name:34s} | target = stock vs fp64: normed {ulp:.3f}" + f" | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 2, f"the drafter runs at tp2, got {WORLD} ranks" + try: + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..2e09105e8fcb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,755 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``, 2 ranks (SM107). + +Its own launcher, as the catalog's collective matrices are: run plainly it spawns +one rank process per device in CUDA_VISIBLE_DEVICES and kills them all on the +first failure or on the deadline. A rank whose peers never push polls forever, +so the deadline is what keeps a broken exchange from wedging the calling run. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_mc_allreduce_norm_op_matrix.py + +Each rank builds the three multicast workspaces (8 / 16 / 48 tokens) the target +builds, over the torch.distributed group (the non-MPI path of +``_make_mnnvl_mcast_buffer``). Every rank can regenerate every rank's inputs from +the shared seed convention, which is what keeps each reference arithmetic rather +than a second collective. Producer pushes (``prepushed`` 2 / 3) are emulated by +storing the canonicalized rows through the multicast mapping at the call parity +the flags select, which is exactly what the producing kernels store. + +Not a pytest module, despite the ``check_*`` bodies: they are one fixed sequence +inside a single 2-rank job whose workspace state (call parities) carries from +one to the next. The collected entry point is +``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +finalize_entry: Any = None +ar_module: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +# ---------------------------------------------------------------- inputs + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + """One rank's [m, H] BF16 rows; ``rank`` -1 for tensors every rank holds alike.""" + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return _rows(seed + 7, -1, m, 4.0) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _finalize_inputs(seed: int, rank: int, m: int, topk: int, dropped: int): + """One rank's unfinalized MoE outputs: fc2 [m * topk, H], perm [m, topk], fp32 scales, shared [m, H].""" + g = _gen(seed + 23, rank) + rows = m * topk + fc2 = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + perm = torch.randperm(rows, generator=g, device="cuda").to(torch.int32) + if dropped: + drop = torch.randperm(rows, generator=g, device="cuda")[:dropped] + perm[drop] = -1 + perm = perm.view(m, topk).contiguous() + scales = torch.rand(m, topk, generator=g, device="cuda", dtype=torch.float32) + shared = _rows(seed + 29, rank, m, 0.5) + return fc2, perm, scales, shared + + +# ---------------------------------------------------------------- workspace / producer emulation + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor, exchange_set: int) -> None: + """Store this rank's rows into every rank's exchange, as a producer kernel does before a prepushed call. + + Slot of (token t, cluster CTA c): set * 2 * half + parity(t, c) * half + (t * world + rank) * H / 2 + c * H / 16 + words, parity read from the CTA's flag; -0.0 halves are stored as +0.0 (0x80000000 is the empty word). + """ + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = ( + exchange_set * 2 * half + + parity.view(m, CLUSTER, 1) * half + + (t * WORLD + RANK) * (HIDDEN // 2) + + c * chunk + + w + ) + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +# ---------------------------------------------------------------- gates + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +#: Share of normed elements one BF16 step off the float64 reference: measured at most 1.1e-4 per call over +#: every call of this matrix on GB300 (see ``_assert_normed``). The gate sits ~18x above that and far below what a +#: wrong statistic produces (most elements move). +NORMED_STEP_FRACTION = 0.002 +NORMED_STEP_SEEN: List[float] = [] + + +def _assert_normed(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """At most one BF16 rounding step from the float64 reference, and rarely that. + + The kernel's RMS statistic is a shuffle / warp / cluster tree of FP32 partial sums and its ``rsqrtf`` is + approximate (~2^-22 relative): the scaled value moves by ~1e-6 relative, which changes the BF16 rounding only + for values within that distance of a rounding midpoint -- one step, never two. The fraction bound catches a + statistic that is systematically off (a missing rank, a wrong H), which moves most elements. + """ + assert out.dtype == ref.dtype == torch.bfloat16 and out.shape == ref.shape, where + o, r = out.float(), ref.float() + _, exponent = torch.frexp(torch.maximum(o.abs(), r.abs())) + step = torch.ldexp(torch.ones_like(o), exponent - 8) + diff = (o - r).abs() + assert bool((diff <= step).all()), f"{where}: max diff {diff.max().item()} beyond one BF16 step" + frac = (diff != 0).float().mean().item() + NORMED_STEP_SEEN.append(frac) + assert frac <= NORMED_STEP_FRACTION, f"{where}: {frac:.4f} of normed elements off by one step" + + +def _assert_quant(outs, ref, m: int, where: str) -> None: + """The MXFP8 copy: bit-exact against the recipe applied to the op's own normed rows (a pure function of + them), and equal to the reference's wherever the reference's normed 32-column block is the op's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + mxfp8_rows, + sf_written, + ) + + q, sf = outs[2], outs[3] + own_q, own_sf = mxfp8_rows(outs[0]) + written = sf_written(m, HIDDEN // 32, sf.device) + _assert_bitwise(q, own_q, f"{where} e4m3 rows") + _assert_bitwise(sf[written], own_sf[written], f"{where} UE8M0 scales") + same_block = (outs[0].view(m, -1, 32) == ref[0].view(m, -1, 32)).all(-1) + qb = q.view(torch.uint8).view(m, -1, 32) + rqb = ref[2].view(torch.uint8).view(m, -1, 32) + assert bool((qb == rqb).all(-1)[same_block].all()), ( + f"{where}: e4m3 differs on a block whose normed rows agree" + ) + + +def _check_outputs(outs, ref, m: int, quant: bool, where: str) -> None: + assert len(outs) == (4 if quant else 2), (where, len(outs)) + _assert_bitwise(outs[1], ref[1], f"{where} updated residual") + _assert_normed(outs[0], ref[0], f"{where} normed") + if quant: + _assert_quant(outs, ref, m, where) + + +# ---------------------------------------------------------------- calls + + +def _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, opts): + return allreduce_entry.m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + ws.max_tokens, + quant, + prepushed, + opts, + ) + + +def _plain_case(spec: Dict[str, Any], seed: int, where: str) -> None: + m, maxt, quant, prepushed = spec["m"], spec["maxt"], spec["quant"], spec["prepushed"] + ws = WORKSPACES[maxt] + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] if spec["x2"] else None + residual, norm_w = _residual(seed, m), _norm_w(seed) + x, x2 = xs[RANK], (None if x2s is None else x2s[RANK]) + if prepushed >= 2: + _push(ws, x, 0) + if prepushed == 3: + _push(ws, x2, 1) + with validating(allreduce_entry.m3_mc_allreduce_norm): + outs = _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, spec["opts"]) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, residual, norm_w, EPS, quant, spec.get("drafter", False) + ) + _check_outputs(outs, ref, m, quant, where) + dist.barrier() + + +# ---------------------------------------------------------------- checks + + +def check_plain_cells() -> None: + """Every cell of ``m3_mc_allreduce_norm.CELLS`` through ``__call__`` under ``validating``, gated on its + ``reference``.""" + for i, cell in enumerate(allreduce_entry.m3_mc_allreduce_norm.CELLS): + _plain_case(cell.spec, 1000 + 17 * i, cell.why) + + +def check_finalize_cells() -> None: + """Every cell of ``m3_mc_finalize_allreduce_norm.CELLS``, gated on its ``reference``.""" + op = finalize_entry.m3_mc_finalize_allreduce_norm + for i, cell in enumerate(op.CELLS): + spec = cell.spec + m, maxt, quant, topk = spec["m"], spec["maxt"], spec["quant"], spec["topk"] + seed = 2000 + 17 * i + ws = WORKSPACES[maxt] + per_rank = [_finalize_inputs(seed, r, m, topk, spec["dropped"]) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + with validating(op): + outs = op( + fc2, + perm, + scales, + shared, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + maxt, + quant, + spec["opts"], + ) + torch.cuda.synchronize() + ref = op.reference( + [p[0] for p in per_rank], + [p[1] for p in per_rank], + [p[2] for p in per_rank], + [p[3] for p in per_rank], + residual, + norm_w, + EPS, + quant, + ) + _check_outputs(outs, ref, m, quant, cell.why) + dist.barrier() + + +def _capture_reference( + updated: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor +) -> None: + """The capture the boundary replaces: m3_spec_tail.capture_norm of the updated residual (no second addend).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_spec_tail, + ) + + m3_spec_tail.capture_norm(updated, None, weight, eps, out) + + +def check_capture_norm_in_boundary() -> None: + """CAPTURE_NORM_IN_BOUNDARY: with cap_out the boundary also writes the capture, BITWISE equal to capture_norm of + its own updated residual, into a strided [M, H] block of a [rows, 3H] capture buffer (other bytes untouched); its + normed / updated / MXFP8 outputs are bitwise the ones of the same call without cap_out. The row kernel (the plain + entry at 2 ranks: x, x + x2, prepushed) and the cluster kernel (MoE finalize), 1 / 4 / 8 rows, residual scales + from 1e-3 to 1e3 (the statistic's range), the fc_norm eps 1e-6 and 1e-5.""" + cap_eps_values = (1e-6, 1e-5) + for case, (m, scale, kind) in enumerate( + ( + (1, 1.0, "plain"), + (4, 1e3, "plain_x2"), + (8, 1e-3, "prepushed"), + (4, 1.0, "prepushed"), + (1, 4.0, "finalize"), + (4, 1e3, "finalize"), + (8, 1e-3, "finalize"), + ) + ): + seed = 5000 + 31 * case + cap_eps = cap_eps_values[case % 2] + ws = WORKSPACES[8] + residual = (_residual(seed, m).float() * scale).to(torch.bfloat16) + norm_w = _norm_w(seed) + cap_w = _norm_w(seed + 101) + quant = case % 3 == 1 + outs_by = [] + for with_cap in (False, True): + buf = torch.full((m + 3, 3 * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + cap = buf[:m, HIDDEN : 2 * HIDDEN] if with_cap else None + if kind == "finalize": + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, m, 4, 1) + fc2 = (fc2.float() * scale).to(torch.bfloat16) + outs = finalize_entry.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), + ws.flags_for(CLUSTER), RANK, WORLD, CLUSTER, False, 8, quant, 7, + cap, cap_w if with_cap else None, cap_eps, + ) # fmt: skip + else: + x = (_rows(seed, RANK, m).float() * scale).to(torch.bfloat16) + x2 = ( + (_rows(seed + 3, RANK, m).float() * scale).to(torch.bfloat16) + if kind == "plain_x2" + else None + ) + prepushed = 2 if kind == "prepushed" else 0 + if prepushed: + _push(ws, x, 0) + outs = allreduce_entry.m3_mc_allreduce_norm( + x, x2, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), ws.flags_for(CLUSTER), RANK, WORLD, + CLUSTER, False, 8, quant, prepushed, (7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + cap, cap_w if with_cap else None, cap_eps, + ) # fmt: skip + torch.cuda.synchronize() + outs_by.append((outs, buf)) + dist.barrier() + (plain, _), (capped, buf) = outs_by + where = f"capture in boundary, {kind}, {m} rows, scale {scale}" + assert len(plain) == len(capped), where + for i, (a, b) in enumerate(zip(plain, capped)): + if i == 3: + # The scale buffer is torch.empty: compare the bytes the kernel writes (the valid rows'). + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written, + ) + + written = sf_written(m, HIDDEN // 32, a.device) + a, b = a[written], b[written] + _assert_bitwise(b, a, f"{where}: output {i} with vs without cap_out") + ref = torch.empty((m, HIDDEN), dtype=torch.bfloat16, device="cuda") + _capture_reference(capped[1], cap_w, cap_eps, ref) + torch.cuda.synchronize() + _assert_bitwise(buf[:m, HIDDEN : 2 * HIDDEN], ref, f"{where}: capture vs capture_norm") + untouched = buf.clone() + untouched[:m, HIDDEN : 2 * HIDDEN] = 7.0 + assert bool((untouched == 7.0).all()), f"{where}: bytes outside the capture block changed" + dist.barrier() + + +def check_the_target_wrappers_select_bucket_and_options() -> None: + """The target's own wrappers (bucket by token count, option table) on top of the entries, every token count of + the 8-token bucket plus one of each larger bucket, and the finalize wrapper; the MXFP8 publication is the + prequant slot's business and is left off.""" + buckets = ar_module.McBuckets([WORKSPACES[b] for b in BUCKETS]) + for m in (*range(1, 9), 11, 40): + seed = 3000 + m + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + normed, updated = ar_module.m3_mc_allreduce_norm( + xs[RANK], residual, norm_w, EPS, buckets, x2=x2s[RANK] + ) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, False) + _check_outputs([normed, updated], ref, m, False, f"wrapper m={m}") + dist.barrier() + for m in (9, 33): + seed = 3100 + m + per_rank = [_finalize_inputs(seed, r, m, 4, 2) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + normed, updated = ar_module.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, buckets + ) + torch.cuda.synchronize() + ref = finalize_entry.m3_mc_finalize_allreduce_norm.reference( + *[[p[j] for p in per_rank] for j in range(4)], residual, norm_w, EPS, False + ) + _check_outputs([normed, updated], ref, m, False, f"finalize wrapper m={m}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """Two boundaries captured in one graph (8- and 16-token buckets), replayed with fresh inputs. + + The decode graph holds every boundary call; each replay rotates the same per-CTA parities the eager calls do, + so a replay that reused a stale half or skipped a re-arm reads the previous replay's rows and fails the + bitwise residual gate (every replay gets new inputs). + """ + shapes = ((3, 8, 15), (12, 16, 15)) + norm_w = _norm_w(4000) + bufs = [] + for m, maxt, opts in shapes: + bufs.append( + dict( + m=m, + maxt=maxt, + opts=opts, + x=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + x2=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + residual=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + ) + + def body(): + return [ + _call_plain( + WORKSPACES[b["maxt"]], b["x"], b["x2"], b["residual"], norm_w, False, 0, b["opts"] + ) + for b in bufs + ] + + def fill(seed: int): + refs = [] + for j, b in enumerate(bufs): + s = seed + 31 * j + xs = [_rows(s, r, b["m"]) for r in range(WORLD)] + x2s = [_rows(s + 3, r, b["m"]) for r in range(WORLD)] + b["x"].copy_(xs[RANK]) + b["x2"].copy_(x2s[RANK]) + b["residual"].copy_(_residual(s, b["m"])) + refs.append( + allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, b["residual"].clone(), norm_w, EPS, False + ) + ) + torch.cuda.synchronize() + dist.barrier() + return refs + + refs = fill(4100) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm = body() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(warm[j], refs[j], b["m"], False, f"graph warm-up m={b['m']}") + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = body() + for rep in range(4): + refs = fill(4200 + 100 * rep) + graph.replay() + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(captured[j], refs[j], b["m"], False, f"replay {rep} m={b['m']}") + dist.barrier() + del graph + dist.barrier() + + +def check_the_gate_sees_a_missing_rank() -> None: + """The residual gate is not blind to the reduction: dropping any one rank's partial from the reference + moves most elements of the updated residual (each rank's rows are distinct by construction).""" + m, seed = 4, 5000 + xs = [_rows(seed, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + full = allreduce_entry.m3_mc_allreduce_norm.reference(xs, None, residual, norm_w, EPS, False) + for r in range(WORLD): + partial = [x if i != r else torch.zeros_like(x) for i, x in enumerate(xs)] + wrong = allreduce_entry.m3_mc_allreduce_norm.reference( + partial, None, residual, norm_w, EPS, False + ) + frac = (wrong[1] != full[1]).float().mean().item() + assert frac > 0.9, f"dropping rank {r} changes only {frac:.3f} of the residual" + dist.barrier() + + +def _expect_rejected(fn, match: str, where: str) -> None: + try: + fn() + except AssertionError as exc: + assert match in str(exc), f"{where}: rejected for the wrong reason: {exc}" + else: + raise AssertionError(f"{where}: the guard admitted it") + + +def check_guards() -> None: + """``is_valid`` refuses what the op accepts and answers wrongly. Nothing here reaches the kernel: the guard + raises first, on every rank alike, so no rank is left polling.""" + m, seed = 4, 6000 + ws8, ws16 = WORKSPACES[8], WORKSPACES[16] + x, residual, norm_w = _rows(seed, RANK, m), _residual(seed, m), _norm_w(seed) + op = allreduce_entry.m3_mc_allreduce_norm + + def call(ws_uc, flags, norm): + with validating(op): + op( + x, + None, + residual, + norm, + EPS, + ws_uc, + ws8.mc.data_ptr(), + flags, + RANK, + WORLD, + CLUSTER, + False, + 8, + False, + 0, + 15, + ) + + strided = torch.empty(2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[::2] + strided.copy_(norm_w) + _expect_rejected( + lambda: call(ws8.uc, ws8.flags_for(CLUSTER), strided), "contiguous", "strided norm_w" + ) + _expect_rejected( + lambda: call(ws8.uc[: ws8.uc.numel() // 2], ws8.flags_for(CLUSTER), norm_w), + "words", + "short workspace", + ) + _expect_rejected( + lambda: call(ws8.uc, ws16.flags_for(CLUSTER), norm_w), "counters", "another bucket's flags" + ) + + fin = finalize_entry.m3_mc_finalize_allreduce_norm + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, 9, 4, 0) + bad = perm.clone() + bad[0, 0] = fc2.shape[0] + + def call_fin(): + with validating(fin): + fin( + fc2, + bad, + scales, + shared, + _residual(seed, 9), + norm_w, + EPS, + ws16.uc, + ws16.mc.data_ptr(), + ws16.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + 16, + False, + 15, + ) + + _expect_rejected(call_fin, "perm spans", "perm past the FC2 rows") + dist.barrier() + + +CHECKS = ( + check_plain_cells, + check_finalize_cells, + check_capture_norm_in_boundary, + check_the_target_wrappers_select_bucket_and_options, + check_cuda_graph_replays, + check_the_gate_sees_a_missing_rank, + check_guards, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, allreduce_entry, finalize_entry, ar_module, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 2, f"the entries are certified at 2 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as _finalize, + ) + + allreduce_entry, finalize_entry, ar_module = _plain, _finalize, _ar + # Spawned by file path, not collected by pytest: reach the sibling helper the plain way (confined to the + # rank processes). /comm -> tests/unittest/_torch/modeling_v2. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peers' polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print( + f"[rank {RANK}] {len(CHECKS)} checks passed; normed one-step fraction max " + f"{max(NORMED_STEP_SEEN):.5f} over {len(NORMED_STEP_SEEN)} calls", + flush=True, + ) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py new file mode 100644 index 000000000000..99df8fb953ef --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification of the MoE engine's all-reduce push (``m3_moe_engine`` with ``arp``), 2 ranks (SM107). + +The decode forward's MoE boundary (MOE_PUSH rows): the engine pushes its bf16 sum (routed + shared) into set A of +the boundary's multicast exchange while it writes it, and the boundary all-reduce runs with ``prepushed`` 2, reading +the pushed words instead of loading and pushing the rows itself. Certified against the unpushed boundary the target +runs otherwise: the same engine call without ``arp``, then ``m3_mc_allreduce_norm`` with ``prepushed`` 0. Every +output must be bit for bit the same: the engine's own output tensor, the updated residual, the normed rows and, with +``quant``, the MXFP8 hand-off it publishes. Both read the same bf16 words, so the reduction and the norm see equal +inputs; any difference is a push that stored the wrong word, slot or parity. + +Each rank holds its own expert shard (operand seed per rank) and the same hidden rows, residual and norm weight, as +the TP2 decode step does. Eager calls alternate pushed and unpushed in both orders so both call parities of the +exchange are exercised; a CUDA graph of the pushed engine + boundary is replayed on fresh rows against eager +unpushed references. The engine's early PDL release (EARLY_AR, the target's default) is checked against the release +after the push: scheduling only, the same bits. + +Its own launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is: run plainly it spawns one rank process per device +in CUDA_VISIBLE_DEVICES and kills them all on the first failure or on the deadline (a rank whose peer never pushes +polls forever). The collected entry point is ``test_modeling_v2_m3_moe_push_op_matrix.py``. + + CUDA_VISIBLE_DEVICES=0,1 python _m3_moe_push_op_matrix.py +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_moe_engine requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1200 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +MAX_TOKENS = 8 +EMPTY_WORD = -2147483648 +RESERVED_SMS = 24 +OPERAND_SEED = 20261013 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +ar = None +sf_written: Any = None +engine_op: Any = None +ar_push_desc: Any = None +moe_state: Any = None +operands: Any = None +prequant: Any = None +RANK = 0 +WORLD = 1 +WS: Any = None + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return (torch.randn(m, HIDDEN, generator=_gen(seed + 7), device="cuda") * 4.0).to( + torch.bfloat16 + ) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _make_workspace(): + """The 8-token bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed + group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * MAX_TOKENS * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=MAX_TOKENS, + ) + ws.flags_for(CLUSTER) + return ws + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +_STATES = {} +_LAYERS = {} + + +def _layer(m: int, case: str, early_release: bool = True): + """This rank's engine layer for ``m`` tokens: its own expert shard; ``concentrated`` routes every token to the + same four experts. ``early_release``: the state's EARLY_AR build (the target's default) or the release after the + push.""" + if (m, early_release) not in _STATES: + sms = torch.cuda.get_device_properties(0).multi_processor_count + _STATES[m, early_release] = moe_state( + torch.device("cuda"), + num_tokens=m, + num_ctas=sms - RESERVED_SMS, + shared=True, + early_release=early_release, + ) + key = (m, case, early_release) + if key not in _LAYERS: + ops = operands.experts(OPERAND_SEED + RANK) + bias = ops.bias.clone() + if case == "concentrated": + bias[list(operands.CONCENTRATED)] += 1.0 + _LAYERS[key] = _STATES[m, early_release].layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=ops.shared, + ) + return _LAYERS[key] + + +def _desc(): + flags = WS.flags_for(CLUSTER) + return ar_push_desc( + WS.mc.data_ptr(), flags.data_ptr(), RANK, WORLD, WS.max_tokens, flags.device + ) + + +def _boundary(x, residual, norm_w, quant, prepushed): + prequant.SLOT.clear() + outs = ar.m3_mc_allreduce_norm( + x, residual, norm_w, EPS, WS, use_pdl=True, quant=quant, prepushed=prepushed + ) + assert outs is not None, "the boundary all-reduce declined its operands" + pub = prequant.SLOT.pending + prequant.SLOT.clear() + return outs, (None if pub is None else (pub[3].clone(), pub[4].clone())) + + +def _unpushed(hidden, layer, residual, norm_w, quant): + out = engine_op(hidden, layer) + (normed, res), pub = _boundary(out, residual.clone(), norm_w, quant, 0) + return out.clone(), normed.clone(), res.clone(), pub + + +def _pushed(hidden, layer, residual, norm_w, quant): + out = engine_op(hidden, layer, arp=_desc()) + (normed, res), pub = _boundary(out, residual.clone(), norm_w, quant, 2) + return out.clone(), normed.clone(), res.clone(), pub + + +def _compare(got, ref, quant: bool, where: str) -> None: + _assert_bitwise(got[0], ref[0], f"{where} engine output") + _assert_bitwise(got[2], ref[2], f"{where} updated residual") + _assert_bitwise(got[1], ref[1], f"{where} normed") + if quant: + assert got[3] is not None and ref[3] is not None, f"{where}: no MXFP8 hand-off published" + _assert_bitwise(got[3][0], ref[3][0], f"{where} e4m3 rows") + # Only the scale atoms of the m rows are written; the swizzled buffer's padding is not. + written = sf_written(got[0].shape[0], HIDDEN // 32, got[3][1].device) + _assert_bitwise(got[3][1][written], ref[3][1][written], f"{where} UE8M0 scales") + + +# ---------------------------------------------------------------- checks + + +def check_pushed_boundary_is_the_unpushed_one() -> None: + """Every token count, both routing cases, with and without the MXFP8 hand-off, both call orders.""" + for m in range(1, MAX_TOKENS + 1): + for case in ("random", "concentrated"): + layer = _layer(m, case) + for quant in (False, True): + for seed in (1, 2): + s = ( + 1000 * m + + 10 * seed + + (5 if quant else 0) + + (500 if case == "concentrated" else 0) + ) + hidden = operands.hidden_states(m, s) + residual, norm_w = _residual(s, m), _norm_w(s) + where = f"m={m} {case} quant={quant} seed={seed}" + if seed == 1: + ref = _unpushed(hidden, layer, residual, norm_w, quant) + got = _pushed(hidden, layer, residual, norm_w, quant) + else: + got = _pushed(hidden, layer, residual, norm_w, quant) + ref = _unpushed(hidden, layer, residual, norm_w, quant) + torch.cuda.synchronize() + _compare(got, ref, quant, where) + + +def check_graph_replays() -> None: + """A captured pushed engine + prepushed boundary, replayed on fresh rows, equals the eager unpushed path.""" + for m in (1, 4, 8): + layer = _layer(m, "random") + hidden = operands.hidden_states(m, 90000 + m).clone() + residual = _residual(90000 + m, m).clone() + norm_w = _norm_w(90000 + m) + desc = _desc() + for _ in range( + 2 + ): # warm both parities eagerly; every push is consumed by its prepushed boundary + _pushed(hidden, layer, residual, norm_w, False) + _unpushed(hidden, layer, residual, norm_w, False) + torch.cuda.synchronize() + dist.barrier() + res_in = residual.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + g_out = engine_op(hidden, layer, arp=desc) + g_normed, g_res = ar.m3_mc_allreduce_norm( + g_out, res_in, norm_w, EPS, WS, use_pdl=True, prepushed=2 + ) + torch.cuda.synchronize() + dist.barrier() + for rep in range(6): + s = 91000 + 100 * m + rep + fresh_h, fresh_r = operands.hidden_states(m, s), _residual(s, m) + hidden.copy_(fresh_h) + res_in.copy_(fresh_r) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + got = (g_out.clone(), g_normed.clone(), g_res.clone(), None) + ref = _unpushed(fresh_h, layer, fresh_r, norm_w, False) + torch.cuda.synchronize() + _compare(got, ref, False, f"graph m={m} replay {rep}") + dist.barrier() + + +def check_early_release_is_scheduling_only() -> None: + """The pushed engine + boundary with the early PDL release and with the release after the push: the same bits.""" + for m in range(1, MAX_TOKENS + 1): + for quant in (False, True): + s = 50000 + 10 * m + (5 if quant else 0) + hidden = operands.hidden_states(m, s) + residual, norm_w = _residual(s, m), _norm_w(s) + early = _pushed(hidden, _layer(m, "random", True), residual, norm_w, quant) + late = _pushed(hidden, _layer(m, "random", False), residual, norm_w, quant) + torch.cuda.synchronize() + _compare(early, late, quant, f"m={m} quant={quant} early vs late release") + + +CHECKS = [ + check_pushed_boundary_is_the_unpushed_one, + check_early_release_is_scheduling_only, + check_graph_replays, +] + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, ar, sf_written, engine_op, ar_push_desc, moe_state, operands, prequant, WS + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 2, f"the push is certified at 2 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_prequant as _pq # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written as _sf_written, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + M3MoeState, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + ar_push_desc as _desc_fn, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as _op, + ) + + ar, prequant, moe_state, ar_push_desc, engine_op = _ar, _pq, M3MoeState, _desc_fn, _op + sf_written = _sf_written + # Spawned by file path, not collected by pytest: the sibling moe/ operands the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "moe")) + import _m3_moe_operands as _operands + + operands = _operands + WS = _make_workspace() + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peer's polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_rank_job.py new file mode 100644 index 000000000000..2d024dc2d3e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_rank_job.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run this target's collective rank job from pytest, and report what it did. + +The same contract as ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``: the +launcher beside this file owns the per-device rank processes and the deadline +that kills a wedged collective; this only selects the devices, starts it in a +fresh interpreter (by file path: this tree is not a package) and turns its exit +code into an assertion. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +WORLD_SIZE = 2 +"""tp2's world size -- the topology these entries are certified for.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the launcher's own deadline, so its message wins the race.""" + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + if len(devices) < WORLD_SIZE: + pytest.skip( + f"this entry is certified at world size {WORLD_SIZE}; only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def _load_launcher_constants(launcher: Path): + """Read the launcher's module-level budget without importing it by name (it imports only torch at + module scope).""" + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str) -> str: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices; returns its stdout.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + + # The ranks import tensorrt_llm absolutely. From a source checkout that is not installed, put the repo root + # (tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root) on their path; from a tests + # tree run against an installed package, they import the installed one. + repo_root = Path(__file__).resolve().parents[8] + if (repo_root / "tensorrt_llm").is_dir(): + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p + ) + else: + assert importlib.util.find_spec("tensorrt_llm") is not None, ( + f"no tensorrt_llm/ at the repo root {repo_root} and no installed tensorrt_llm" + ) + + constants = _load_launcher_constants(launcher) + timeout = constants.DEADLINE_S + _LAUNCHER_GRACE_S + + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) + return completed.stdout diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_allreduce.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_allreduce.py new file mode 100644 index 000000000000..2e3d69a51316 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_allreduce.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the ``m3_allreduce`` certification matrix. + +The matrix is ``_m3_allreduce_op_matrix.py`` beside this file, its own 2-rank launcher; see ``_rank_job``. It drives +every cell of the entry against its reference, then requires the fused draft loop's boundaries built on it (plain +all-reduce, then the fused add + RMSNorm) to equal stock Eagle3's bit for bit under every strategy, and reports both +against an fp64 reference as a warning, one line per strategy and case. +""" + +import warnings + +import _rank_job +import pytest +import torch + +pytestmark = [ + pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _rank_job.WORLD_SIZE, + reason=f"certified at world size {_rank_job.WORLD_SIZE}", + ), +] + +_RESULT_TAG = "[ulp]" + + +# Spawns 2 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_allreduce_op_matrix() -> None: + out = _rank_job.run("m3_allreduce") + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + "draft boundary (= stock Eagle3) vs fp64, max BF16 ulps\n" + "\n".join(lines), stacklevel=1 + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..75f31cdb994b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the multicast all-reduce + RMSNorm entries' certification matrix. + +Covers both ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm`` (one +kernel, one exchange protocol; the finalize entry's reference has its own test in +``test_modeling_v2_m3_mc_finalize_allreduce_norm.py``). The matrix is ``_m3_mc_allreduce_norm_op_matrix.py`` +beside this file, its own 2-rank launcher; see ``_rank_job`` for why. +""" + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + + +# Spawns 2 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_mc_allreduce_norm_op_matrix() -> None: + _rank_job.run("m3_mc_allreduce_norm") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..984ed13ced0e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``m3_mc_finalize_allreduce_norm``: its reference's rounding. The entry's GPU cells (9..48 rows, every one gated on +that reference) run in the multicast all-reduce matrix, ``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``, which +launches both entries of the one kernel at 2 ranks.""" + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + finalize_partial, + m3_mc_finalize_allreduce_norm, +) + + +def test_finalize_partial_rounds_the_routed_sum_before_the_shared_expert() -> None: + """As stock: the finalized routed output is BF16, and the shared expert is added to it (shared.add_(routed)). + Four routed rows of 1.0 at FP32 routing weight 0.2509765625 sum to 1.00390625, which rounds to BF16 1.0; plus the + shared -1 that is 0, where adding before the rounding would leave 0.00390625.""" + fc2 = torch.ones((4, 8), dtype=torch.bfloat16) + perm = torch.arange(4, dtype=torch.int32).view(1, 4) + scales = torch.full((1, 4), 0.2509765625, dtype=torch.float32) + shared = torch.full((1, 8), -1.0, dtype=torch.bfloat16) + out = finalize_partial(fc2, perm, scales, shared) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(out, shared.clone().add_(routed)) + assert torch.equal(out, torch.zeros(1, 8, dtype=torch.bfloat16)) + + +def test_the_cells_cover_every_row_the_target_sends_here() -> None: + """The decode layer's MC finalize (MC_FINALIZE_MIN_TOKENS 9 .. the b32 plan's 32 rows) and the generic layer's + (past 32, up to GENERIC_MC_FINALIZE_MAX_TOKENS 48): every one of those row counts is a certified cell.""" + rows = {cell.spec["m"] for cell in m3_mc_finalize_allreduce_norm.CELLS} + assert set(range(9, 49)) <= rows diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_moe_push_op_matrix.py new file mode 100644 index 000000000000..c548a51e3060 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_moe_push_op_matrix.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the MoE engine's all-reduce push certification (``_m3_moe_push_op_matrix.py`` beside +this file, its own 2-rank launcher; see ``_rank_job``).""" + +import _rank_job +import pytest +import torch + +pytestmark = [ + pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _rank_job.WORLD_SIZE, + reason=f"certified at world size {_rank_job.WORLD_SIZE}", + ), +] + + +# Spawns 2 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_moe_push_op_matrix() -> None: + _rank_job.run("m3_moe_push") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..b77aabe64e20 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_embedding entry of the MiniMax-M3 tp2 target: an exact gather with zero rows outside the +table.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.embedding.m3_embedding import ( # noqa: E501 + m3_embedding as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" + +VOCAB, HIDDEN = 4096, 6144 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + ids = torch.randint( + 0, VOCAB, (spec["rows"],), generator=g, device="cuda", dtype=spec["id_dtype"] + ) + if spec["pad"]: + ids[1::3] = -1 + with validating(op): + out = op(ids, table) + assert torch.equal(out, op.reference(ids, table)) + if spec["pad"]: + assert not out[1::3].any(), "padded ids must embed to zero rows" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py new file mode 100644 index 000000000000..e44ad56d743d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" + +import torch + + +def swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): + """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" + v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) + return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..b455f6f2a049 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_bf16_linear entry of the MiniMax-M3 tp2 target (the drafter's projections).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_bf16_linear import ( # noqa: E501 + m3_bf16_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the + row, RMS well under one; measured on GB300: worst 0.5 unit.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w = (torch.randn(spec["n"], spec["k"], generator=g, device="cuda") * 0.02).to(torch.bfloat16) + with validating(op): + out = op(x, w) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(x, w), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py new file mode 100644 index 000000000000..3af452daa1d9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the dense MLP's activation seam (``modeling._dense_swiglu``, the ``m3_dense_swiglu`` C++ op): between +the dense MLP's two decode GEMVs, the seam's down-GEMV output is bit for bit the Triton swigluoai's +(``m3_silu_and_mul``) at every decode row count, eager and in a CUDA graph, with PDL; and the seam declines (to the +Triton kernel) what it does not cover.""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.mx_gemv import ( # noqa: E501 + op as gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + DENSE_GATE_UP, +) + +assert torch.cuda.is_available(), "the dense seam test requires a CUDA device" + +ALPHA, LIMIT, BETA = 1.702, 7.0, 1.0 +ROWS = list(range(1, modeling.TINY_MAX_TOKENS + 1)) +CORE = SimpleNamespace(swiglu_alpha=ALPHA, swiglu_limit=LIMIT, swiglu_beta=BETA) + + +def _weights(n, k, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return wq.view(n, k), wsf + + +GATE_UP = None +DOWN = None + + +@pytest.fixture(scope="module", autouse=True) +def _mlp_weights(): + global GATE_UP, DOWN + GATE_UP = _weights(*DENSE_GATE_UP, 11) + DOWN = _weights(*DENSE_DOWN, 13) + + +def _gemv(x, wts): + out = gemv.m3_mx_gemv_bf16(x.contiguous(), *wts, use_pdl=True) + assert out is not None, ( + "the target's bf16 GEMV serves every decode row of the dense projections" + ) + return out + + +def _mlp(x, seam: bool): + """gate_up GEMV, the activation (the seam, or the Triton kernel), down GEMV.""" + gu = _gemv(x, GATE_UP) + if seam: + act = modeling._dense_swiglu(CORE, gu) + else: + act = modeling.m3_silu_and_mul(gu, LIMIT, ALPHA, BETA) + return _gemv(act, DOWN) + + +def _x(rows, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(rows, DENSE_GATE_UP[1], generator=g, device="cuda") * 2).to(torch.bfloat16) + + +@pytest.fixture +def op_calls(monkeypatch): + """Records each call that reaches the C++ op through the seam.""" + calls = [] + real = modeling.m3_dense_swiglu + + class _Recording: + CELLS = real.CELLS + + def __call__(self, gu, *args): + calls.append(gu.shape[0]) + return real(gu, *args) + + monkeypatch.setattr(modeling, "m3_dense_swiglu", _Recording()) + return calls + + +@pytest.mark.parametrize("rows", ROWS) +def test_seam_on_is_bitwise_seam_off(op_calls, rows) -> None: + for trial in range(3): + x = _x(rows, 100 * rows + trial) + on = _mlp(x, True) + off = _mlp(x, False) + torch.cuda.synchronize() + assert torch.equal(on.view(torch.int16), off.view(torch.int16)), f"trial {trial}" + assert op_calls == [rows] * 3, "the seam runs the C++ op at every decode row count" + + +@pytest.mark.parametrize("rows", [1, 4, modeling.TINY_MAX_TOKENS]) +def test_seam_in_a_graph_is_bitwise_seam_off(rows) -> None: + """Both paths captured; replays on new inputs (PDL edges inside the graph).""" + x = _x(rows, 7) + _mlp(x, True) # compile outside the capture + graph_on = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph_on): + on = _mlp(x, True) + _mlp(x, False) + graph_off = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph_off): + off = _mlp(x, False) + for step in range(3): + x.copy_(_x(rows, 31 + step)) + graph_on.replay() + graph_off.replay() + torch.cuda.synchronize() + assert torch.equal(on.view(torch.int16), off.view(torch.int16)), f"replay {step}" + + +def test_seam_declines_what_it_does_not_cover(op_calls, monkeypatch) -> None: + """Rows past the op's certified rows and other activation parameters take the Triton kernel (the same bits).""" + max_rows = max(cell.spec["rows"] for cell in modeling.m3_dense_swiglu.CELLS) + # Past the decode rows the GEMV declines too; the seam only reads the gate_up rows, so any BF16 rows will do. + g = torch.Generator(device="cuda").manual_seed(5) + gu = torch.randn(max_rows + 1, DENSE_GATE_UP[0], generator=g, device="cuda").to(torch.bfloat16) + got = modeling._dense_swiglu(CORE, gu) + assert op_calls == [], "rows past the op's certified rows" + assert torch.equal(got, modeling.m3_silu_and_mul(gu, LIMIT, ALPHA, BETA)) + other = SimpleNamespace(swiglu_alpha=1.0, swiglu_limit=LIMIT, swiglu_beta=BETA) + gu = _gemv(_x(4, 5), GATE_UP) + got = modeling._dense_swiglu(other, gu) + assert op_calls == [], "other activation parameters" + assert torch.equal(got, modeling.m3_silu_and_mul(gu, LIMIT, 1.0, BETA)) + + +def test_seam_needs_no_capture_scope(op_calls, monkeypatch) -> None: + """The op replaces only the activation kernel: the small-batch decode graphs are captured outside the FlashInfer + MXFP8 capture scope (r4.10.1 never engaged at 1-8 rows).""" + monkeypatch.setattr(modeling, "flashinfer_mxfp8_decode_graph_capture_active", lambda: False) + modeling._dense_swiglu(CORE, _gemv(_x(4, 5), GATE_UP)) + assert op_calls == [4], "outside the FlashInfer capture scope" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py new file mode 100644 index 000000000000..2b8b3a4f2e0c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise oracle of the dense MLP's swigluoai C++ op (ops/gemm/m3_dense_swiglu, +trtllm::minimax_m3_sm107_tp2_swiglu_oai) against the Triton silu_and_mul kernel the dense MLP runs today +(modules/swiglu.py at limit 7, alpha 1.702, beta 1): every bf16 gate bit pattern against special and random up values, +every up bit pattern against special and random gates, random pairs, and the dense MLP's shapes at every decode row +count, eager and replayed from a CUDA graph, with and without the programmatic launch.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.silu_and_mul +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_dense_swiglu import ( # noqa: E501 + ALPHA, + BETA, + LIMIT, + m3_dense_swiglu, +) + +assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_swiglu_oai"): + pytest.skip( + "needs the TP2 op library with minimax_m3_sm107_tp2_swiglu_oai", allow_module_level=True + ) + +# The TP2 dense MLP's activation width (gate_up 3072 -> d 1536). +DENSE_D = modeling.DENSE_DOWN[1] + + +def _triton(x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=LIMIT, swiglu_alpha=ALPHA, swiglu_beta=BETA + ) + + +def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" + return out + + +def _row(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + return torch.cat([gate.view(1, -1), up.view(1, -1)], dim=1).contiguous() + + +def _all_bf16() -> torch.Tensor: + return ( + torch.arange(1 << 16, dtype=torch.int32, device="cuda").to(torch.int16).view(torch.bfloat16) + ) + + +def _specials(n: int, seed: int) -> list: + g = torch.Generator(device="cuda").manual_seed(seed) + vals = [0.0, -0.0, 1.0, -1.0, 6.99, 7.0, 7.01, -7.0, -7.01, 1e-38, -1e-38, 1e-40, 88.0, -88.0, 3e38, -3e38, + float("inf"), float("-inf"), float("nan")] # fmt: skip + rows = [torch.full((n,), v, dtype=torch.bfloat16, device="cuda") for v in vals] + for scale in (0.05, 1.0, 4.0, 30.0): + rows.append((scale * torch.randn(n, generator=g, device="cuda")).to(torch.bfloat16)) + rows.append( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + return rows + + +def _assert_bitwise(got: torch.Tensor, ref: torch.Tensor, what: str) -> None: + bad = (got.view(torch.int16) != ref.view(torch.int16)).nonzero().flatten() + assert bad.numel() == 0, f"{what}: {bad.numel()} mismatches, first at {bad[:4].tolist()}" + + +def test_every_gate_pattern_is_bitwise_the_triton_kernel() -> None: + gate = _all_bf16() + for j, up in enumerate(_specials(gate.numel(), 1)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"up set {j}") + + +def test_every_up_pattern_is_bitwise_the_triton_kernel() -> None: + up = _all_bf16() + for j, gate in enumerate(_specials(up.numel(), 2)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"gate set {j}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: + g = torch.Generator(device="cuda").manual_seed(100 + seed) + n = 1 << 22 + + def pattern() -> torch.Tensor: + bits = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + return bits.to(torch.int16).view(torch.bfloat16) + + x = _row(pattern(), pattern()) + _assert_bitwise(_op(x), _triton(x), f"seed {seed}") + + +@pytest.mark.parametrize("pdl", [True, False], ids=["pdl", "plain"]) +@pytest.mark.parametrize("rows", range(1, modeling.TINY_MAX_TOKENS + 1)) +def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pdl: bool) -> None: + """The TP2 dense MLP's gate_up rows (a GEMV output's range plus values past the clamps), eager and replayed from a + CUDA graph after fresh rows are written in place, against the Triton kernel on the same rows.""" + g = torch.Generator(device="cuda").manual_seed(2000 + rows) + x = (3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, DENSE_D : DENSE_D + 8] = x[:, :8] + _assert_bitwise(_op(x, pdl), _triton(x), f"eager rows {rows}") + out = torch.empty(rows, DENSE_D, dtype=torch.bfloat16, device="cuda") + _op(x, pdl) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out.copy_(_op(x, pdl)) + x.copy_((3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(out, _triton(x), f"graph rows {rows}") + + +def test_the_op_declines_other_parameters_and_layouts() -> None: + x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + + +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: + """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton + kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" + rows, d = cell.spec["rows"], cell.spec["d"] + g = torch.Generator(device="cuda").manual_seed(3000 + rows) + x = (3.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, d : d + 8] = x[:, :8] + _assert_bitwise(_op(x), _triton(x), cell.why) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py new file mode 100644 index 000000000000..b44554a8d717 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_draft_lowm entry: the Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M +kernels (direct SIMT and split-K tensor-core), at the TP2 drafter shards, at the rows its table certifies. + +Every cell is held to the fp64 reference within the derived band. The drafter's projection call +(``M3Eagle3Draft.linear``) at a covered row count is the entry's call bit for bit, at an uncovered one the in-tree +GEMM's. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + TACTICS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 + m3_draft_lowm as op, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" + +# Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). +# +# Kernel and reference multiply the same BF16 operands (every product exact in fp32); the kernel sums in fp32 (the +# direct kernel one dot product in K order, the split-K kernel tensor-core blocks and then its split partials), the +# reference in fp64, and both round once to BF16. With w ~ N(0, 1 / K) and x ~ N(0, 1) the outputs are O(1) and the +# fp32 order difference is ~sqrt(K) * 2^-24 of the summed magnitudes, a few 1e-6 at K = 18432: far below half a BF16 +# ulp, so the two roundings disagree by at most one ulp of the element -- up to two units of the row scale for an +# element in its row maximum's binade -- and only where the sum sits that close to a rounding boundary, a few in ten +# thousand elements, hence 0.1 ulp RMS. The same derivation as the decode GEMV's band (test_modeling_v2_m3_mx_gemv), +# with no slack under it: a second ulp would mean a wrong sum. +ELEMENT_ULP = 2.0 +RMS_ULP = 0.1 + +#: The drafter's projections (N, K) and their weight keys on ``M3Eagle3Draft.w``. +DRAFTER_LINEARS = { + (4608, 12288): "qkv", + (6144, 4096): "o", + (18432, 6144): "gu", + (6144, 9216): "dn", + (6144, 18432): "fc", +} + + +@pytest.fixture(autouse=True) +def _policy(): + modeling.set_require(False) + modeling.reset_logged() + yield + modeling.set_require(False) + modeling.reset_logged() + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + return x, w + + +def test_the_table_is_the_tp2_drafter() -> None: + assert set(TACTICS) == set(DRAFTER_LINEARS) + assert {c.spec["nk"] for c in op.CELLS} == set(TACTICS) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + kind, tactic = tactic_for(n, k, m) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + args = ( + x, + w, + out, + kind, + tactic, + True, + ) # the target launches with PDL when the environment enables it + certifying(op) + assert op(*args) is out + ref = op.reference(*args) + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=RMS_ULP) + # PDL only changes the launch attribute. + plain = torch.empty_like(out) + op(x, w, plain, kind, tactic, False) + assert torch.equal(plain.view(torch.int16), out.view(torch.int16)) + + +def _draft(linears=DRAFTER_LINEARS): + """A drafter holding only its projection weights and the low-M bookkeeping (the module is not built).""" + draft = object.__new__(modeling.M3Eagle3Draft) + draft.w = {key: build(*nk, 1, seed=nk[0] + nk[1])[1] for nk, key in linears.items()} + draft._lowm_built, draft._lowm_engaged = set(), set() + draft._push_arm, draft._pushed = None, 0 + return draft + + +@pytest.fixture(scope="module") +def prebuilt(): + draft = _draft() + assert draft.prebuild_linears() == len(op.CELLS) + assert draft.prebuild_linears() == 0, "a second prebuild finds nothing left to compile" + return draft + + +@pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) +def test_the_projection_is_the_entry_and_falls_back_past_its_rows(prebuilt, nk) -> None: + n, k = nk + key = DRAFTER_LINEARS[nk] + w = prebuilt.w[key] + use_pdl = get_env_enable_pdl() + modeling.set_require(True) + # The certified rows the target enables (DRAFT_LOWM_MAX_ROWS caps the table: the batch-4 / batch-8 cells are off + # by default since they read flat in-model). + for m in sorted(modeling.DRAFT_LOWM_ROWS[nk]): + assert modeling.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt.linear(x, key) + kind, tactic = tactic_for(n, k, m) + want = op(x, w, torch.empty(m, n, dtype=x.dtype, device="cuda"), kind, tactic, use_pdl) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)), m + # Rows the table does not certify (the b2 rows, and past it) or that DRAFT_LOWM_MAX_ROWS leaves off: the in-tree + # GEMM, silently, even under require. + for m in sorted( + (set(range(1, modeling.TINY_MAX_TOKENS + 2)) | set(ROWS[nk])) + - set(modeling.DRAFT_LOWM_ROWS[nk]) + ): + assert not modeling.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt.linear(x, key) + assert torch.equal(got.view(torch.int16), F.linear(x, w).view(torch.int16)), m + + +def test_a_projection_outside_the_table_is_the_in_tree_gemm() -> None: + draft = _draft({(256, 512): "o"}) + x, _ = build(256, 512, 1, seed=3) + assert draft._lowm_linear(draft.w["o"], x) is None + got = draft.linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_a_covered_call_not_prebuilt_raises_under_require() -> None: + draft = _draft({(6144, 4096): "o"}) + x, _ = build(6144, 4096, 1, seed=5) + modeling.set_require(True) + with pytest.raises(modeling.FastPathUnavailable, match="not compiled before warmup"): + draft.linear(x, "o") + # Outside require it falls back to the in-tree GEMM. + modeling.set_require(False) + got = draft.linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_guard_refuses_a_strided_activation() -> None: + x, w = build(6144, 4096, 2, seed=6) + wide = torch.zeros(2, 8192, dtype=torch.bfloat16, device="cuda") + wide[:, :4096] = x + kind, tactic = tactic_for(6144, 4096, 2) + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op( + wide[:, :4096], + w, + torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda"), + kind, + tactic, + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..ebfee3b50ee0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp32_linear entry of the MiniMax-M3 tp2 target (the router beyond the decode routers' rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """torch's default fp32 matmul precision on this hardware may be TF32 (~2^-11 relative per product): the band is + stated in fp32 ulps of each row's largest logit and covers it (K = 6144 TF32 products: ~2^13 ulps worst element); + measured on GB300 with the engine's defaults: worst element ~1700 ulps, RMS ~300.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(spec["e"], spec["k"], generator=g, device="cuda") * 0.02 + with validating(op): + out = op(x, w) + assert out.dtype == torch.float32 and out.shape == (spec["m"], spec["e"]) + assert_within_ulp(out, op.reference(x, w), element_ulp=8192.0, rms_ulp=2048.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..9057207107af --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv entry: the decode MXFP8 GEMV on a producer-quantized activation. + +Operands as the target holds them: the weight through mxfp8_quantize with the swizzled scale layout (the MXFP8 linear +method's storage), the activation through mxfp8_quantize with the swizzled layout (what the boundary all-reduce and +the attention merge publish). +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" + + +# Element band, in ULP units of the row scale: see compare. +ELEMENT_ULP = 1.5 + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Element ELEMENT_ULP of the row scale, relative RMS 0.1 ulp. + + Kernel and reference multiply bit-identical dequantized operands; the kernel sums in fp32 (tensor-core blocks, + then the split-K partials), the reference in fp64. Over K <= 6144 that order difference is ~1e-6 relative, so the + two bf16 roundings disagree only where the sum sits within that of a rounding boundary: one bf16 ulp of that + element, which for an element in its row maximum's binade is up to two of ``assert_within_ulp``'s units of the row + scale, and a few in ten thousand elements in aggregate. Any fp32 sum shows it: torch's IEEE fp32 matmul of the + same dequantized operands measured at most 1.22 units against the fp64 reference over every row count 1..48 on + the o_proj and the QKV (6 seeds each, GB300), the kernel at most 1.25, both at most 0.009 ulp RMS. The band + leaves the margin between that and a full extra ulp. The kernel is the same on SM107 (the tp2 table only picks + split and ring), and its TP2 projections' K is 1536..6144. + """ + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=0.1) + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + return x, xq.view(m, k), xsf, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + _, xq, xsf, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(xq, xsf, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + compare(out, op.reference(xq, xsf, w, wsf)) + + +def test_an_unserved_shape_returns_none() -> None: + """Shapes outside the tuned table fall back at the call site; the op says so rather than computing.""" + _, xq, xsf, w, wsf = build(256, 512, 1, seed=1) + assert op(xq, xsf, w, wsf) is None + + +def test_guard_refuses_linear_activation_scales() -> None: + """A linear [M, K / 32] scale buffer is read as the swizzled one.""" + x, _, _, w, wsf = build(6144, 4096, 4, seed=2) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, False) + with pytest.raises(AssertionError, match="swizzled"): + with validating(op): + op(xq.view(4, 4096), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + _, xq, xsf, w, wsf = build(6144, 4096, 1, seed=3) + wide = torch.zeros(6144, 8192, dtype=w.dtype, device="cuda") + wide[:, :4096] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :4096], wsf) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..4055f2d8db4e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv_bf16 entry: the decode MXFP8 GEMV quantizing a bf16 activation in its prologue. + +The weight as the MXFP8 linear method stores it (mxfp8_quantize, swizzled scales). Besides the reference band, the +prologue's quantization is held to mxfp8_quantize itself: the GEMV on a bf16 x must equal, bit for bit, the +prequantized GEMV on mxfp8_quantize(x) -- the claim the target relies on when it mixes the two paths. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv_bf16 import ( # noqa: E501 + m3_mx_gemv_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" + +# The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. +ELEMENT_ULP = 1.5 + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return x, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(x, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + # The band of test_modeling_v2_m3_mx_gemv.compare: the quantized operands are the reference's exactly (checked + # below), the sums differ only in order. + assert_within_ulp(out, op.reference(x, w, wsf), element_ulp=ELEMENT_ULP, rms_ulp=0.1) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + pre = m3_mx_gemv(xq.view(m, k), xsf, w, wsf, use_pdl=True) + assert pre is not None and torch.equal(out.view(torch.int16), pre.view(torch.int16)), ( + "the prologue's MXFP8 quantization differs from mxfp8_quantize" + ) + + +def test_a_strided_activation_is_declined() -> None: + """The prologue reads x densely; the op declines a strided x (None) rather than computing on it.""" + x, w, wsf = build(6144, 4096, 2, seed=1) + wide = torch.zeros(2, 8192, dtype=torch.bfloat16, device="cuda") + wide[:, :4096] = x + assert op(wide[:, :4096], w, wsf) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..f59ca4d385b4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_gemm entry of the MiniMax-M3 tp2 target (the decode-graph GEMM on a prequantized +activation).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most + 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, False), + element_ulp=2.0, + rms_ulp=0.25, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..5d272004698e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_linear entry of the MiniMax-M3 tp2 target: MXFP8 quantize + GEMM on each backend the +engine selects (native CUTLASS, FlashInfer, the per-bucket tuned quantize + GEMM).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" + +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and + one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" + spec = cell.spec + if ( + spec["path"] == "flashinfer" + and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None + ): + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + args = (x, w, w_sf, alpha, torch.bfloat16, *_FLAGS[spec["path"]]) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py new file mode 100644 index 000000000000..e685bc5df277 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for MX-TAC: the pinned decode-graph MXFP8 decision (CuTe DSL quantize + CuTe DSL GEMM) gives the same bits +as every other candidate stock's AutoTuner can pick (TRT-LLM quantize, CUTLASS GEMM), the same on every call.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.modeling import ( # noqa: E501 + _pinned_mxfp8_gemm, + _pinned_mxfp8_quantize, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + +assert torch.cuda.is_available(), "MX-TAC requires a CUDA device" + +# The TP2 projections at decode rows the GEMVs do not serve (12, 16) and a few beyond. +CASES = [ + ((4992, 6144), 16), + ((6144, 4096), 12), + ((3072, 6144), 24), + ((12288, 6144), 32), + ((6144, 1536), 48), +] + + +@pytest.mark.parametrize("nk,m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_pinned_decision_matches_every_stock_candidate(nk, m, scale) -> None: + n, k = nk + g = torch.Generator(device="cuda").manual_seed(n + k + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = (torch.randn(m, k, generator=g, device="cuda") * scale).to(torch.bfloat16) + x[0, ::32] *= 1e3 # one outlier per 32-element block on the first row + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + wq = wq.view(n, k) + xq, xsf = _pinned_mxfp8_quantize(x) + tq, tsf = torch.ops.trtllm.mxfp8_quantize(x, True) + assert torch.equal(xq.view(torch.uint8).view(m, k), tq.view(torch.uint8).view(m, k)) + assert torch.equal(xsf.view(torch.uint8).flatten(), tsf.view(torch.uint8).flatten()) + xq = xq.view(m, k) + out = _pinned_mxfp8_gemm(xq, xsf, wq, wsf, torch.bfloat16) + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + for tactic in (FlashInferMXFP8GemmRunner.CUTE_DSL, FlashInferMXFP8GemmRunner.CUTLASS): + ref = runner.forward([xq, xsf, wq, wsf], tactic=tactic) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), tactic + again = _pinned_mxfp8_gemm(xq, xsf, wq, wsf, torch.bfloat16) + assert torch.equal(out.view(torch.int16), again.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py new file mode 100644 index 000000000000..a514ebbe4b38 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the uGPU target head (ops/gemm/m3_ugpu_head): the die-placed copy of the head shard is the weight bit +for bit with its halves in the two locality domains; the module's own weight (what every other reader of the head +uses) is never placed; the per-die split GEMM gives the in-tree head's logits bit for bit at the certified row counts, +eager and in a CUDA graph; an in-place change of the weight makes the copy stale.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm import ( # noqa: E501 + m3_ugpu_head as op, +) +from tensorrt_llm._torch.modules.linear import Linear + +assert torch.cuda.is_available(), "the uGPU head requires a CUDA device" + + +def _locality(): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_locality import ( # noqa: E501 + LocalityMap, + ) + + try: + return LocalityMap.get(torch.cuda.current_device()) + except RuntimeError as err: # one locality domain (not a two-die GPU) + pytest.skip(f"no two locality domains: {err}") + + +@pytest.fixture(scope="module") +def placed(): + lm = _locality() + head = Linear(op.HIDDEN, op.TARGET_SHARD, bias=False, dtype=torch.bfloat16) + head.to("cuda") + g = torch.Generator(device="cuda").manual_seed(20261015) + with torch.no_grad(): + head.weight.copy_( + (0.02 * torch.randn(op.TARGET_SHARD, op.HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ) + ) + before = head.weight.detach().clone() + x = torch.randn(16, op.HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + stock = {m: head.apply_linear(x[:m], None) for m in (4, 8, 12, 16)} + ugpu = op.prepare_head(head) + assert ugpu is not None and ugpu.source is head.weight and op.current(ugpu, head.weight) + return lm, head, ugpu, before, x, stock + + +def test_no_reader_of_the_module_sees_placed_memory(placed) -> None: + """A kernel reading die-placed memory from all SMs runs 1.5-1.7x slower: every tensor of the head module (what + the in-tree GEMM, prefill logits, warmup and fallbacks read) must stay in ordinary memory; only the op's copy is + placed.""" + lm, head, ugpu, before, _, _ = placed + for name, t in list(head.named_parameters()) + list(head.named_buffers()): + end = t.data_ptr() + t.numel() * t.element_size() - 1 + assert lm.domain_of(t.data_ptr()) == -1 and lm.domain_of(end) == -1, name + assert t.data_ptr() != ugpu.weight.data_ptr(), name + assert torch.equal(head.weight, before) + + +def test_placement_is_exact_and_per_die(placed) -> None: + lm, head, ugpu, before, _, _ = placed + w = ugpu.weight + assert torch.equal(w, before) + row = w.shape[1] * w.element_size() + assert lm.domain_of(w.data_ptr()) == 0 + assert lm.domain_of(w.data_ptr() + ugpu.split * row - 1) == 0 + assert lm.domain_of(w.data_ptr() + ugpu.split * row) == 1 + assert lm.domain_of(w.data_ptr() + w.numel() * w.element_size() - 1) == 1 + assert (ugpu.split * row) % lm.granularity == 0 and abs(2 * ugpu.split - w.shape[0]) < 2 * 512 + + +@pytest.mark.parametrize( + "cell", op.m3_ugpu_head.CELLS, ids=[c.why[:44] for c in op.m3_ugpu_head.CELLS] +) +def test_split_head_is_the_in_tree_head_bit_for_bit(placed, cell) -> None: + m = cell.spec["m"] + _, head, ugpu, _, x, stock = placed + got = op.m3_ugpu_head(x[:m], ugpu) + torch.cuda.synchronize() + assert torch.equal(got, stock[m]) + assert torch.equal(got, op.m3_ugpu_head.reference(x[:m], ugpu)) + assert torch.equal( + head.apply_linear(x[:m], None), stock[m] + ) # stock readers of the placed weight + + +@pytest.mark.parametrize( + "cell", op.m3_ugpu_head.CELLS, ids=[c.why[:44] for c in op.m3_ugpu_head.CELLS] +) +def test_split_head_in_a_cuda_graph(placed, cell) -> None: + m = cell.spec["m"] + _, _, ugpu, _, x, stock = placed + xs = x[:m].clone() + box = {} + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + box["y"] = op.m3_ugpu_head(xs, ugpu) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + box["y"] = op.m3_ugpu_head(xs, ugpu) + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + assert torch.equal(box["y"], stock[m]) + + +def test_an_in_place_weight_change_makes_the_copy_stale(placed) -> None: + _, head, ugpu, _, _, _ = placed + assert op.current(ugpu, head.weight) + with torch.no_grad(): + head.weight.add_(0) # an in-place write bumps the version counter + assert not op.current(ugpu, head.weight) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py new file mode 100644 index 000000000000..e7a31d3ae752 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rule 1b gate a for the TP2 decode GEMV order override (kernels/mx_gemv/op.STRIDED_K_TILES: strided k-tiles): +for every overridden (projection, rows) the output is the same on 3 eager calls and on a CUDA graph replay +(deterministic), and its max abs / max rel error vs the fp64 product of the dequantized operands stays within the +shipped order's band (the shared table's config on the same operands, over several seeds).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.mx_gemv import ( # noqa: E501 + op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + mxfp8_dequant, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + +CASES = [(nk, m) for nk in (QKV_SPARSE, QKV, O_PROJ, DENSE_GATE_UP) for m in (1, 2, 4, 8)] + + +def _err(y, ref): + d = (y.double() - ref).abs() + return float(d.max()), float((d / ref.abs().amax(dim=1, keepdim=True).clamp_min(1e-30)).max()) + + +@pytest.mark.parametrize("nk, m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +def test_gate_a(nk, m) -> None: + n, k = nk + cfg = op.get_config(n, k, m) + assert cfg is not None and len(cfg) > 2 and cfg[2] == 0, ( + "the override is off or does not cover this call" + ) + shipped = cfg[:2] # the contiguous order + worst_new, worst_ship = [0.0, 0.0], [0.0, 0.0] + for seed in range(4): + g = torch.Generator(device="cuda").manual_seed(1000 * seed + m) + w = (torch.randn(n, k, device="cuda", generator=g) * 2).to(torch.float8_e4m3fn) + ws = torch.randint(118, 124, (n * k // 32,), dtype=torch.uint8, device="cuda", generator=g) + x = (torch.randn(m, k, device="cuda", generator=g) * (0.5 + seed)).to(torch.bfloat16) + q, s = torch.ops.trtllm.mxfp8_quantize(x, True) + ref = ( + mxfp8_dequant(q.view(torch.uint8), s).double() + @ mxfp8_dequant(w.view(torch.uint8), ws).double().t() + ) + ys = [m3_mx_gemv(q, s, w, ws, True, cfg, True, None) for _ in range(3)] + assert all(y is not None for y in ys), "the GEMV declined the override" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + out = [] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out.append(m3_mx_gemv(q, s, w, ws, True, cfg, True, None)) + graph.replay() + torch.cuda.synchronize() + for y in ys[1:] + [out[0]]: + assert torch.equal(ys[0].view(torch.int16), y.view(torch.int16)), ( + "not deterministic (eager / graph)" + ) + ship = m3_mx_gemv(q, s, w, ws, True, shipped, True, None) + torch.cuda.synchronize() + for acc, y in ((worst_new, ys[0]), (worst_ship, ship)): + e = _err(y, ref) + acc[0], acc[1] = max(acc[0], e[0]), max(acc[1], e[1]) + print( + f"\nR1B1 tp2 {n}x{k} m={m} cfg {cfg}: max abs / rel vs fp64 {worst_new[0]:.3e} / {worst_new[1]:.3e} | shipped " + f"order {worst_ship[0]:.3e} / {worst_ship[1]:.3e}", + flush=True, + ) + # The band: the shipped order's own error (bf16 output rounding dominates), with 25 % headroom. + assert worst_new[0] <= 1.25 * worst_ship[0] and worst_new[1] <= 1.25 * worst_ship[1] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..b5cb1d2e9109 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_moe_operands.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One tp2 rank's MiniMax-M3 MoE layer operands for the decode-engine entries' tests. + +The routed experts are random checkpoint-format NVFP4 tensors ([w3 (up); w1 (gate)] rows, e4m3 block scales) put +through TRT-LLM's own TRTLLM-Gen NVFP4 loader shuffle (``NVFP4TRTLLMGenFusedMoEMethod._shuffle_all_experts``), which +writes the buffers the engines read in place. The per-expert scalars follow the loader's conventions (act_beta and +act_clamp divided by fc31_alpha, fc31_scale_c = fc2_input_scale x fc31_alpha) and are sized so the FC1 sums land +around the SwiGLU-OAI clamp and the outputs near unit scale. The shared expert's MXFP8 weights come from +mxfp8_quantize with the swizzled scale layout, as the MXFP8 linear method stores them. +""" + +import functools +from types import SimpleNamespace + +import torch + +# Hidden size, the routed experts' intermediate per rank at TP2 (3072 / 2: every expert on every rank), experts, +# top-k, and the shared expert's intermediate per rank. +H, I_TP, NUM_EXPERTS, TOP_K = 6144, 1536, 128, 4 +SHARED_I = 1536 +SWIGLU_ALPHA, SWIGLU_LIMIT = 1.702, 7.0 +ROUTED_SCALING = 2.0 +MINIMAX2_ROUTING = 5 # RoutingMethodType.MiniMax2 +# NVFP4 global scales: 448 x 6 over the amax each side is sized for. +FC31_INPUT_SCALE = 448.0 * 6.0 / 8.0 +FC2_INPUT_SCALE = 448.0 * 6.0 / 24.0 + + +def _rand_e4m3_scales(shape, gen) -> torch.Tensor: + """E4M3 block scales in [0.5, 2): biased exponents 6 and 7, any mantissa.""" + return torch.randint(0x30, 0x40, shape, dtype=torch.uint8, device="cuda", generator=gen).view( + torch.float8_e4m3fn + ) + + +@functools.lru_cache(maxsize=None) +def experts(seed: int = 20261013): + """The routed experts and the router of one MoE layer, in the layout and conventions of the TRTLLM-Gen loader. + + Returns a SimpleNamespace: ``tensors`` (the ten expert operands in ``M3MoeState.layer``'s order), ``bias``, + ``router_weight``, ``fc31_input_scale`` and ``shared`` (gate_up, its scales, down, its scales). + """ + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + gen = torch.Generator(device="cuda").manual_seed(seed) + e = NUM_EXPERTS + w31 = torch.randint( + 0, 256, (e, 2 * I_TP, H // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + w2 = torch.randint(0, 256, (e, H, I_TP // 2), dtype=torch.uint8, device="cuda", generator=gen) + w31s = _rand_e4m3_scales((e, 2 * I_TP, H // 16), gen) + w2s = _rand_e4m3_scales((e, H, I_TP // 16), gen) + module = SimpleNamespace( + w3_w1_weight=w31, + w2_weight=w2, + w3_w1_weight_scale=w31s, + w2_weight_scale=w2s, + is_gated_activation=True, + bias=False, + ) + NVFP4TRTLLMGenFusedMoEMethod()._shuffle_all_experts(module) + torch.cuda.synchronize() + + # Raw FC1 sums: rms(x) x FC31_INPUT_SCALE x rms(E2M1 value) x rms(block scale) x sqrt(H) ~ 1.0e5 at rms(x) = 1; + # fc31_alpha brings them to ~1.5 in the dequantized domain, the FC2 alpha the outputs to ~1. + base = 1.5 / (FC31_INPUT_SCALE * 2.93 * 1.3 * H**0.5) + fc31_alpha = base * (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) + fc2_alpha = (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) / ( + FC2_INPUT_SCALE * 2.93 * 1.3 * I_TP**0.5 + ) + tensors = ( + w31, + w31s, + w2, + w2s, + fc31_alpha, + FC2_INPUT_SCALE * fc31_alpha, # fc31_scale_c + torch.full((e,), SWIGLU_ALPHA, device="cuda"), # act_alpha + 1.0 / fc31_alpha, # act_beta: the (up + 1) offset over fc31_alpha + SWIGLU_LIMIT / fc31_alpha, # act_clamp over fc31_alpha + fc2_alpha, + ) + bias = (0.02 * torch.randn(e, device="cuda", generator=gen)).float() + router_weight = (torch.randn(e, H, device="cuda", generator=gen) / H**0.5).float().contiguous() + fc31_input_scale = torch.tensor([FC31_INPUT_SCALE], dtype=torch.float32, device="cuda") + + gu = (0.02 * torch.randn(2 * SHARED_I, H, device="cuda", generator=gen)).to(torch.bfloat16) + dn = (0.05 * torch.randn(H, SHARED_I, device="cuda", generator=gen)).to(torch.bfloat16) + gu_q, gu_sf = torch.ops.trtllm.mxfp8_quantize(gu, True) + dn_q, dn_sf = torch.ops.trtllm.mxfp8_quantize(dn, True) + shared = (gu_q.view(2 * SHARED_I, H), gu_sf, dn_q.view(H, SHARED_I), dn_sf) + return SimpleNamespace( + tensors=tensors, + bias=bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared=shared, + ) + + +# Routing cases: random (up to 4 M distinct experts), and every token on the same four experts (the bias outweighs +# any sigmoid difference), which leaves later slot groups empty. +CONCENTRATED = (5, 40, 77, 126) + + +def layer_for(state, case: str, shared: bool = True): + """One layer on ``state`` over :func:`experts` for a routing case. ``shared=False``: the shared expert's weights + are zero, so the engine returns the routed part alone (bf16(0 + bf16(routed))).""" + ops = experts() + bias = ops.bias.clone() + if case == "concentrated": + bias[list(CONCENTRATED)] += 1.0 + shared_weights = ops.shared + if not shared: + gu, gu_sf, dn, dn_sf = ops.shared + shared_weights = (torch.zeros_like(gu), gu_sf, torch.zeros_like(dn), dn_sf) + return state.layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=shared_weights, + ) + + +def stock_routed( + hidden: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """The stock TRTLLM-Gen NVFP4 MoE (fp4_quantize, then fp4_block_scale_moe_runner with SwiGLU-OAI) on the given + routing, over :func:`experts`: what the built-in MiniMaxM3MoE computes for the routed part.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + + ops = experts() + w31, w31s, w2, w2s, fc31_alpha, scale_c, act_alpha, act_beta, act_clamp, fc2_alpha = ops.tensors + xq, xsf = torch.ops.trtllm.fp4_quantize(hidden, ops.fc31_input_scale, 16, False, False) + out = torch.ops.trtllm.fp4_block_scale_moe_runner( + None, None, xq, xsf.view(torch.float8_e4m3fn).flatten(), + w31, w31s, None, act_alpha, act_beta, act_clamp, w2, w2s, None, + scale_c, fc31_alpha, fc2_alpha, + NUM_EXPERTS, TOP_K, None, None, I_TP, 0, NUM_EXPERTS, ROUTED_SCALING, MINIMAX2_ROUTING, True, 0, + topk_weights=topk_weights.to(torch.bfloat16), topk_ids=topk_ids.to(torch.int32), + ) # fmt: skip + return out[0] if isinstance(out, (list, tuple)) else out + + +def hidden_states(m: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(m, H, device="cuda", generator=gen).to(torch.bfloat16) + + +def check_routed_part(op, state, m: int, seed: int) -> None: + """Exact, independent of the reference: the engine's routing (its route_dbg record: ids, then the bf16 weights' + fp32 bits) is the reference routing bit for bit, and its routed part is the stock TRTLLM-Gen MoE's output on that + routing bit for bit -- the claim the engines are built on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=False) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + stock = stock_routed(hidden, ids, weights) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)), ( + "routed part differs from TRTLLM-Gen's" + ) + + +def stock_shared(hidden: torch.Tensor) -> torch.Tensor: + """The stock MXFP8 shared expert as the built-in GatedMLP runs it (MXFP8LinearMethod's native path): mxfp8_quantize + + mxfp8_mxfp8_gemm, the fused swiglu (limit 7, alpha 1.702, beta 1), mxfp8_quantize + mxfp8_mxfp8_gemm.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + from tensorrt_llm._torch.modules.swiglu import swiglu + + gu, gu_sf, dn, dn_sf = experts().shared + one = torch.ones(1, dtype=torch.float32, device=hidden.device) + a, s = torch.ops.trtllm.mxfp8_quantize(hidden.contiguous(), True) + y = torch.ops.trtllm.mxfp8_mxfp8_gemm(a, s, gu, gu_sf, one, torch.bfloat16) + h = swiglu(y, swiglu_limit=SWIGLU_LIMIT, swiglu_alpha=SWIGLU_ALPHA, swiglu_beta=1.0) + a2, s2 = torch.ops.trtllm.mxfp8_quantize(h.contiguous(), True) + return torch.ops.trtllm.mxfp8_mxfp8_gemm(a2, s2, dn, dn_sf, one, torch.bfloat16) + + +def check_full_output(op, state, m: int, seed: int) -> None: + """Exact: the engine's bf16(shared + routed) is the stock MoE block's shared_output.add_(routed_output) bit for bit + (stock TRTLLM-Gen routed part on the engine's routing, stock MXFP8 shared expert).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=True) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + ref = stock_shared(hidden).add_(stock_routed(hidden, ids, weights)) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), ( + "bf16(shared + routed) differs from stock" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py new file mode 100644 index 000000000000..6d8fe90ad9d7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared operands and gates of the router entries' tests (m3_router_gemv, m3_router_gemv_fp4, m3_router_tc: one +kernel source).""" + +import zlib + +import torch + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv import ( # noqa: E501 + m3_router_gemv as gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv import ( # noqa: E501 + near_tie, + nvfp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 + m3_router_gemv_fp4 as gemv_fp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__all__ = [ + "SF_SCALE", + "_build", + "_compare_fp4", + "_compare_logits", + "_seed", + "gemv", + "gemv_fp4", + "tc", +] + +#: The experts' NVFP4 input scale the tc cells use: 448 * 6 / 8, block scales well inside E4M3. +SF_SCALE = 336.0 + + +def _build(m: int, e: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(e, k, generator=g, device="cuda") * 0.02 + return x, w + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +def _compare_logits(out: torch.Tensor, ref: torch.Tensor) -> None: + """fp32 logits that differ from a true fp32 product only in summation order. + + Over K = 6144 products a reordered fp32 sum moves by a few ulps of the sum of magnitudes, which near + cancellation is ~sqrt(K) times the row's largest logit's ulp at most. Gated in fp32 ulps of each row's largest + logit: 64 element-wise, 8 in relative RMS. Measured on GB300: GEMV 1.6 / 1.4 ulp, tensor-core router 3.8 / 2.0 + ulp at 2..48 tokens; the tensor-core logits against a true fp32 product (the TF32 rounding the band must see) + sit at 2200+ / 1700 ulp. + """ + assert out.dtype == ref.dtype == torch.float32 and out.shape == ref.shape + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=8.0) + + +def _codes(packed: torch.Tensor) -> torch.Tensor: + """[M, K / 2] packed e2m1 bytes -> [M, K] 4-bit codes, low nibble first.""" + lo, hi = packed & 0xF, packed >> 4 + return torch.stack((lo, hi), dim=-1).reshape(packed.shape[0], -1) + + +def _compare_fp4( + q: torch.Tensor, sf: torch.Tensor, x: torch.Tensor, sf_scale: torch.Tensor +) -> None: + """Scales bit-exact; codes bit-exact except within ``TIE_WINDOW`` of an e2m1 midpoint, where the kernel's + approximate reciprocals choose the side (and then only by one code).""" + ref_q, ref_sf, scaled = nvfp4_quantize(x, sf_scale) + assert q.dtype == torch.uint8 and q.shape == ref_q.shape + assert torch.equal(sf, ref_sf), f"{(sf != ref_sf).sum().item()} block scales differ" + got, want = _codes(q), _codes(ref_q) + differ = got != want + assert not bool((differ & ~near_tie(scaled)).any()), ( + f"{(differ & ~near_tie(scaled)).sum().item()} e2m1 codes differ away from a rounding tie" + ) + magnitude = ((got & 7).int() - (want & 7).int()).abs() + assert bool((magnitude[differ] <= 1).all()), "a tie resolved more than one code away" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py new file mode 100644 index 000000000000..f4116a9c71b2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router of an unfused MoE past the tensor-core router's 48 rows: the fp32 router (``m3_fp32_linear``) is bit for +bit the built-in gate's generation-step arithmetic there, ``F.linear(h.to(float32), w)`` (the gate's TF32 fixed tile +is off in generation steps, its tensor-core router covers 9..16 rows, and below 2048 rows its fp32 cast is the plain +``.to(torch.float32)``), so a 49..192-row decode step routes as the agg canon does.""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear, +) +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 107, reason="SM107 (Rubin) only" +) + +HIDDEN, EXPERTS = 6144, 128 + + +@pytest.mark.parametrize("rows", [49, 64, 128, 192]) +def test_the_fp32_router_is_the_gates_generation_arithmetic(rows: int) -> None: + assert not modeling.ROUTER_TC_GEN.covers(rows), "the tensor-core router stops at 48 rows" + g = torch.Generator(device="cuda").manual_seed(rows) + h = torch.randn(rows, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + w = torch.randn(EXPERTS, HIDDEN, generator=g, device="cuda") / HIDDEN**0.5 + got = m3_fp32_linear(h, w) + want = F.linear(h.to(torch.float32), w) + assert got.dtype == torch.float32 + assert torch.equal(got.view(torch.int32), want.view(torch.int32)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..aad0252e6662 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe_engine entry: the 1..8-token MoE engine on its caller-owned state. + +One Rubin (SM107) rank of the tp2 deployment: 128 experts (all local, routed and shared intermediate 1536 per rank), +top 4, hidden 6144, the shared expert folded in, ``num_ctas`` = SMs - 24 as the target builds it (212 - 24 = 188). +Operands: ``_m3_moe_operands``. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + HIDDEN, + INTER, + SHARED_INTER, + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import ( # noqa: E402 + I_TP, + SHARED_I, + H, + check_full_output, + check_routed_part, + hidden_states, + layer_for, +) +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +_STATES = {} + + +def test_the_operands_are_the_tp2_engine_geometry() -> None: + assert (H, I_TP, SHARED_I) == (HIDDEN, INTER, SHARED_INTER) == (6144, 1536, 1536) + + +def test_the_engine_refuses_the_tp4_width() -> None: + with pytest.raises(ValueError, match="1536"): + M3MoeState(torch.device("cuda"), num_tokens=1, inter=768) + + +def test_the_engine_refuses_more_routed_ctas_than_fc2_blocks() -> None: + """A routed CTA owning none of the 192 FC2 output blocks would deadlock the hand-off (routed-only, 212 SMs).""" + with pytest.raises(ValueError, match="192"): + M3MoeState(torch.device("cuda"), num_tokens=4, num_ctas=193, shared=False) + + +def _state(m: int) -> M3MoeState: + if m not in _STATES: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + _STATES[m] = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + return _STATES[m] + + +def test_a_module_built_with_other_trace_config_compiles_its_own_kernel(monkeypatch) -> None: + """The compile cache keys on the config the kernel module was built with: a state whose module carries an extra + trace-time option (``timing``, which only adds per-CTA stamps) compiles its own kernel, never the plain build's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_moe import ( # noqa: E501 + op as kop, + ) + + hidden = hidden_states(1, seed=5) + op(hidden, layer_for(_state(1), "random")) # the plain build is compiled + n0 = len(kop._compiled) + build = kop._kernel_module + monkeypatch.setattr( + kop, "_kernel_module", lambda config, path: build({**config, "timing": 1}, path) + ) + sms = torch.cuda.get_device_properties(0).multi_processor_count + state = M3MoeState(torch.device("cuda"), num_tokens=1, num_ctas=sms - 24, shared=True) + op(hidden, layer_for(state, "random")) + assert state.mod.TIMING and not _state(1).mod.TIMING + assert len(kop._compiled) == n0 + 1, "the timing build reused another build's compiled kernel" + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """The engine against the plain-torch MoE: BAND_ELEMENT bf16 ulp of the row's largest magnitude per element, + BAND_RMS ulp relative RMS. + + Both sides quantize the same operands with the same recipes (the activation's and the intermediate's NVFP4 codes + and block scales, the shared expert's MXFP8); they differ in summation order (tensor-core fp32 against fp64) and in + the PTX approximations (rcp / ex2 / tanh). Such a difference changes the output only where it moves a value across + a rounding boundary: one bf16 ulp of a slot output or of the sum, or, rarely, one intermediate NVFP4 code, which + reaches the output as one weight column's share. Measured over every cell, both routing cases and 3 seeds each + on the tp4 engine (GB300, intermediate 768): at most 1.48 ulp element, 0.18 ulp RMS; the tp2 engine runs the + same recipes over twice the intermediate. The band leaves ~3x on both; a misread recipe (the activation's + 1 / (sf x (1 / sc)) spelling of the reciprocal on ~0.06% of its codes) measured 12 ulp element, 7 ulp RMS. + """ + assert_within_ulp(out, ref, element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + # The workspace is tagged per call and never reset: a repeated call is the same call. + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + compare(out, op.reference(hidden, layer)) + + +@pytest.mark.parametrize("m", (1, 2, 4, 8)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +@pytest.mark.parametrize("seed", (1, 2, 3)) +@pytest.mark.parametrize("m", (4, 8)) +def test_full_output_is_the_stock_block_bit_for_bit(m, seed) -> None: + """The TP2 operating points (one / two requests verifying 3 drafts): the whole layer output, shared expert + included, against the stock block (TRTLLM-Gen routed + MXFP8 GatedMLP shared, shared.add_(routed)).""" + check_full_output(op, _state(m), m, seed=200 + 10 * m + seed) + + +def test_layers_on_one_state_do_not_leak_into_each_other() -> None: + """Two layers share one state's workspace, in stream order: A, B, A gives A's output twice, bit for bit, and the + caller's ``out`` receives the same rows as a fresh tensor.""" + state = _state(4) + a, b = layer_for(state, "random"), layer_for(state, "concentrated") + hidden = hidden_states(4, seed=7) + first = op(hidden, a) + op(hidden, b) + out = torch.empty(4, H, dtype=torch.bfloat16, device="cuda") + second = op(hidden, a, out=out) + assert second.data_ptr() == out.data_ptr() + assert torch.equal(first.view(torch.int16), out.view(torch.int16)) + + +@pytest.mark.parametrize( + "bad", + ("misaligned hidden", "fp32 out", "short out", "long out", "int32 arp"), +) +def test_malformed_buffers_are_rejected_by_the_call(bad) -> None: + """Why the entry has no is_valid: the compiled call checks each buffer against its signature and raises.""" + state = _state(2) + layer = layer_for(state, "random") + hidden, kw = hidden_states(2, seed=3), {} + if bad == "misaligned hidden": + hidden = torch.randn(2 * H + 8, device="cuda").to(torch.bfloat16)[2 : 2 + 2 * H].view(2, H) + elif bad == "fp32 out": + kw["out"] = torch.empty(2, H, dtype=torch.float32, device="cuda") + elif bad == "short out": + kw["out"] = torch.empty(1, H, dtype=torch.bfloat16, device="cuda") + elif bad == "long out": + kw["out"] = torch.empty(3, H, dtype=torch.bfloat16, device="cuda") + else: + kw["arp"] = torch.zeros(8, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + op(hidden, layer, **kw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py new file mode 100644 index 000000000000..59559dc01d89 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the m3_moe engine's per-call epoch wraps past 2^31 without changing a bit. + +The engine tags its hand-off records with the call epoch (masked low bits plus a marker bit) and counts the B2IMG group +counters up by CNT_PER_CALL per call, waiting for ``(epoch + 1) * CNT_PER_CALL`` in wrapping int32 arithmetic. A state +seeded at epoch 0x7FFFFFF0 (counters seeded consistently, as after that many calls) runs 32 calls across the int32 +wrap; every output equals a fresh state's, bit for bit. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import hidden_states, layer_for # noqa: E402 + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +EPOCH0 = 0x7FFFFFF0 +CALLS = 32 + + +def _int32(v: int) -> int: + v &= 0xFFFFFFFF + return v - (1 << 32) if v >= 1 << 31 else v + + +def _seed(state: M3MoeState, epoch: int) -> None: + """The state as after ``epoch`` calls: every CTA's epoch, and each B2IMG group counter at epoch x CNT_PER_CALL.""" + state.epochs.fill_(_int32(epoch)) + mod = state.mod + if mod.B2IMG: + value = _int32(epoch * mod.CNT_PER_CALL) + for g in range(mod.NG_CAP): + state.cnt[8 * g] = value + torch.cuda.synchronize() + + +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_the_epoch_wraps_bit_for_bit(m) -> None: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + fresh = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + wrapped = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + _seed(wrapped, EPOCH0) + lf, lw = layer_for(fresh, "random"), layer_for(wrapped, "random") + for i in range(CALLS): + x = hidden_states(m, 4000 + i) + a, b = op(x, lf), op(x, lw) + torch.cuda.synchronize() + assert torch.equal(a, b), f"call {i} (epoch {EPOCH0 + i:#x})" + assert int(wrapped.epochs[0]) == _int32(EPOCH0 + CALLS) < 0 # crossed 2^31 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py new file mode 100644 index 000000000000..eb0f3beec25c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Who finalizes the wide target's MoE at each generic row count: past the decode rows, up to +GENERIC_MC_FINALIZE_MAX_TOKENS, the finalize (+ the shared expert) folds into the boundary's multicast all-reduce +(M3DecoderLayer._mc_finalize_boundary); from MOE_TRITON_FINALIZE_MIN_TOKENS, kernels/m3_moe_finalize finalizes it. +The two ranges meet with no gap and no overlap, and the fold declines from shapes alone past its cap. + +CPU only: nothing here launches a kernel. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_mc_allreduce, +) + +H = 6144 + + +def _fold_cap() -> int: + return min(dispatch.GENERIC_MC_FINALIZE_MAX_TOKENS, m3_mc_allreduce.FINALIZE_MAX_TOKENS) + + +def test_the_fold_stays_inside_the_finalize_kernel() -> None: + assert dispatch.GENERIC_MC_FINALIZE_MAX_TOKENS <= m3_mc_allreduce.FINALIZE_MAX_TOKENS + + +def test_the_fold_and_the_triton_finalize_meet() -> None: + """With the fold on (MC_AR_TP2), the Triton finalize starts on the first row past it; with it off, nothing folds.""" + if not dispatch.MC_AR_TP2: + assert dispatch.GENERIC_MC_FINALIZE_MAX_TOKENS == 0 + return + assert dispatch.DECODE_MAX_TOKENS < _fold_cap() + if dispatch.MOE_TRITON_FINALIZE: + assert dispatch.MOE_TRITON_FINALIZE_MIN_TOKENS == _fold_cap() + 1 + + +def _target(): + """A DecodeTarget whose bucket lookup and MC finalize must not be reached (the row gate declines first).""" + + def unreachable(*args, **kwargs): + raise AssertionError("the generic MC finalize ran past its row gate") + + return SimpleNamespace(_bucket_flags=unreachable, _mc_finalize=unreachable) + + +@pytest.mark.parametrize("past", [1, 2, 16]) +def test_the_fold_declines_past_its_cap(past: int) -> None: + hidden = torch.empty((_fold_cap() + past, H), dtype=torch.bfloat16) + assert dispatch.DecodeTarget._generic_mc_finalize(_target(), 3, hidden, hidden, None) is None + + +@pytest.mark.parametrize("rows", [1, 8, 32]) +def test_the_fold_leaves_the_decode_layers_rows(rows: int) -> None: + rows = min(rows, dispatch.DECODE_MAX_TOKENS) + hidden = torch.empty((rows, H), dtype=torch.bfloat16) + assert dispatch.DecodeTarget._generic_mc_finalize(_target(), 3, hidden, hidden, None) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py new file mode 100644 index 000000000000..cd34e710ecfd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The TP2 MoE engine's FLAT stack is applied only up to FLAT_MAX_TOKENS (4) tokens per call. + +At 8 tokens FLAT was +2.96 us/call slower than the per-m tables, so above the cap the engine keeps the per-m +configuration (the engine with the cap at 0, bitwise). CPU only: M3MoeState is constructed with its device queries +stubbed and stopped at the kernel build, so the test reads the engine config and the engagement lines without a GPU. +""" + +import types + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_moe import ( # noqa: E501 + op, +) + + +class _Built(Exception): + """Raised by the stubbed kernel build: the config is final by then.""" + + +def _engine(monkeypatch, num_tokens): + """(config, info lines) of M3MoeState(num_tokens) up to its kernel build.""" + lines = [] + seen = {} + + def build(config, path): + seen["config"] = dict(config) + raise _Built() + + monkeypatch.setattr(op, "is_supported", lambda device=None: (True, "")) + monkeypatch.setattr(op, "_kernel_module", build) + monkeypatch.setattr(op.torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr( + op.torch.cuda, + "get_device_properties", + lambda device: types.SimpleNamespace(multi_processor_count=op.HIDDEN // 32), + ) + from tensorrt_llm.logger import logger + + monkeypatch.setattr(logger, "info", lambda msg, *a, **k: lines.append(str(msg))) + with pytest.raises(_Built): + op.M3MoeState("cuda", num_tokens, num_ctas=op.HIDDEN // 32) + return seen["config"], [line for line in lines if "MiniMax-M3 MoE engine FLAT" in line] + + +def _has_flat(config): + return all(config.get(k) == v for k, v in op.FLAT_STACK.items()) + + +@pytest.mark.parametrize("num_tokens", range(1, 9)) +def test_flat_applies_only_up_to_the_cap(monkeypatch, num_tokens): + config, lines = _engine(monkeypatch, num_tokens) + assert op.FLAT_MAX_TOKENS == 4 + if num_tokens <= 4: + assert _has_flat(config), config + assert lines == [ + "MiniMax-M3 MoE engine FLAT: b2img=0 lg_repl=16 ar_pre=0 late_trigger=0 " + f"tokens={num_tokens}" + ] + else: + # above the cap: the per-m tables, i.e. exactly the cap-0 engine + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 0) + off_config, _ = _engine(monkeypatch, num_tokens) + assert config == off_config + assert lines == [ + f"MiniMax-M3 MoE engine FLAT off above 4 tokens (per-m tables): tokens={num_tokens}" + ] + + +@pytest.mark.parametrize("num_tokens", range(1, 9)) +def test_cap_zero_runs_the_per_m_tables(monkeypatch, num_tokens): + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 0) + config, lines = _engine(monkeypatch, num_tokens) + assert lines == [ + f"MiniMax-M3 MoE engine FLAT off above 0 tokens (per-m tables): tokens={num_tokens}" + ] + assert config.get("b2img") == op.B2IMG_BY_M[num_tokens - 1] + assert config.get("ar_pre") == op.AR_PRE_BY_M[num_tokens - 1] + + +@pytest.mark.parametrize("num_tokens", range(1, 9)) +def test_cap_of_eight_runs_flat_at_every_token_count(monkeypatch, num_tokens): + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 8) + config, lines = _engine(monkeypatch, num_tokens) + assert _has_flat(config), config + assert len(lines) == 1 and lines[0].endswith(f"tokens={num_tokens}") + + +def test_flat_applies_helper(monkeypatch): + """op.flat_applies reads FLAT_MAX_TOKENS per call.""" + assert [op.flat_applies(m) for m in range(1, 9)] == [True] * 4 + [False] * 4 + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 8) + assert all(op.flat_applies(m) for m in range(1, 9)) + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 0) + assert not any(op.flat_applies(m) for m in range(1, 9)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py new file mode 100644 index 000000000000..13cec78a3a07 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: SH_TAKE (the shared-expert CTAs take the extra FC2 output blocks) changes no output bit. + +With 188 routed CTAs, 192 FC2 output blocks leave four routed CTAs with two blocks. Under SH_TAKE the last four +shared-expert CTAs run, after their shared work, one of those blocks each as FC2-only CTAs, with the same tiles, +scales and finalize. ``SH_TAKE_BY_M`` turns it on per token count (shipped all off). For every token count the +engine's output under the gated table (on at 6 / 7 / 8 tokens) equals its output under the shipped table, bit for bit, +eager and replayed from a CUDA graph, for spread and concentrated routing. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_moe import ( # noqa: E501 + op as engine_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as run, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import hidden_states, layer_for # noqa: E402 + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +SHIPPED = engine_op.SH_TAKE_BY_M +GATED = (0, 0, 0, 0, 0, 1, 1, 1) + + +def _state(monkeypatch, m, on): + monkeypatch.setattr(engine_op, "SH_TAKE_BY_M", GATED if on else SHIPPED) + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + return M3MoeState(dev, num_tokens=m, num_ctas=sms - engine_op.SHARED_CTAS, shared=True) + + +@pytest.mark.parametrize("m", range(1, 9)) +def test_tables(monkeypatch, m) -> None: + assert not any(SHIPPED) + on = _state(monkeypatch, m, True) + off = _state(monkeypatch, m, False) + nheavy = (engine_op.HIDDEN // 32) % on.num_ctas + expect = bool(GATED[m - 1]) and 0 < nheavy <= engine_op.SHARED_CTAS + assert bool(on.config.get("sh_take")) == expect + assert "sh_take" not in off.config + + +@pytest.mark.parametrize("case", ["random", "concentrated"]) +@pytest.mark.parametrize("m", range(1, 9)) +def test_bitwise_against_the_shipped_table(monkeypatch, m, case) -> None: + states = {on: _state(monkeypatch, m, on) for on in (True, False)} + layers = {on: layer_for(st, case) for on, st in states.items()} + for i in range(8): + x = hidden_states(m, 6000 + i) + a, b = run(x, layers[True]), run(x, layers[False]) + torch.cuda.synchronize() + assert torch.equal(a, b), f"eager call {i}" + x = hidden_states(m, 6100) + outs = {} + for on in (True, False): + run(x, layers[on]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = run(x, layers[on]) + graph.replay() + torch.cuda.synchronize() + outs[on] = y.clone() + assert torch.equal(outs[True], outs[False]), "graph replay" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py new file mode 100644 index 000000000000..cc92290a9162 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""kernels/m3_moe_finalize: the routed experts' finalize (+ the shared expert) from the TRTLLM-Gen runner's unfinalized +parts, bitwise to the runner's finalizeKernelVecLoad order (acc = 0; acc += w_k * fc2[idx_k] for k in order, skipping +-1; one BF16 rounding) and to ``shared.add_(routed)``, at the rows the wide target routes there (49..16K), with dropped +(-1) slots, padded FC2 rows, eager and CUDA-graph replay. (The runner itself was compared at 512 / 2K / 4K / 16K rows on +SM107, maximum difference 0.)""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_moe_finalize import ( # noqa: E501 + m3_moe_finalize, +) + +H, TOP_K = 6144, 4 + + +def _parts(rows: int, seed: int, pad: int = 128): + g = torch.Generator(device="cuda").manual_seed(seed) + n = rows * TOP_K + fc2 = torch.randn((n + pad, H), device="cuda", generator=g).to(torch.bfloat16) + idx = torch.randperm(n, device="cuda", generator=g).to(torch.int32).reshape(rows, TOP_K) + drop = torch.rand((rows, TOP_K), device="cuda", generator=g) < 0.02 + idx = torch.where(drop, torch.full_like(idx, -1), idx) + w = torch.rand((rows, TOP_K), device="cuda", generator=g).to(torch.bfloat16) + shared = torch.randn((rows, H), device="cuda", generator=g).to(torch.bfloat16) + return fc2, w, idx.contiguous(), shared + + +def _reference(fc2, w, idx): + acc = torch.zeros((w.shape[0], H), dtype=torch.float32, device=fc2.device) + for k in range(TOP_K): + valid = (idx[:, k] >= 0).unsqueeze(1) + x = fc2[idx[:, k].clamp(min=0).long()].float() + acc = torch.where(valid, acc + w[:, k : k + 1].float() * x, acc) + return acc.to(torch.bfloat16) + + +@pytest.mark.parametrize("rows", [49, 64, 65, 1000, 4096, 16384]) +def test_bitwise_to_the_runner_order(rows: int) -> None: + fc2, w, idx, shared = _parts(rows, rows) + ref = _reference(fc2, w, idx) + out = m3_moe_finalize(fc2, w, idx) + assert out is not None and torch.equal(out.view(torch.int16), ref.view(torch.int16)) + fused = m3_moe_finalize(fc2, w, idx, shared) + assert torch.equal(fused.view(torch.int16), shared.clone().add_(ref).view(torch.int16)) + + +def test_graph_replay_matches_eager() -> None: + fc2, w, idx, shared = _parts(512, 7) + eager = m3_moe_finalize(fc2, w, idx, shared) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = m3_moe_finalize(fc2, w, idx, shared) + for seed in (8, 9): + f2, w2, i2, s2 = _parts(512, seed) + fc2.copy_(f2) + w.copy_(w2) + idx.copy_(i2) + shared.copy_(s2) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, m3_moe_finalize(fc2, w, idx, shared)) + assert eager.shape == out.shape + + +def test_declines_other_layouts() -> None: + fc2, w, idx, shared = _parts(64, 3) + assert m3_moe_finalize(fc2.float(), w, idx) is None + assert m3_moe_finalize(fc2, w, idx.long()) is None + assert m3_moe_finalize(fc2, w, idx, shared[:, :128].contiguous()) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..61bef8aacee1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_tc entry of the MiniMax-M3 tp2 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + + +@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +def test_tc_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) + sf_scale = torch.tensor([SF_SCALE], dtype=torch.float32, device="cuda") if spec["fp4"] else None + with validating(tc): + outs = tc(x, w, sf_scale, False) + ref = tc.reference(x, w, sf_scale, False) + assert len(outs) == len(ref) == (3 if spec["fp4"] else 1) + _compare_logits(outs[0], ref[0]) + if spec["fp4"]: + _compare_fp4(outs[1], outs[2], x, sf_scale) + + +def test_tc_reference_is_tf32_and_the_gate_sees_it() -> None: + """The tensor-core router rounds the weight to TF32: against a true fp32 product its logits sit far outside + the band, so the TF32 rounding in ``reference`` is load-bearing rather than decoration.""" + x, w = _build(16, 128, 6144, 5) + out = tc(x, w, None, False)[0] + with pytest.raises(AssertionError): + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +def test_guard_refuses_a_misaligned_activation() -> None: + """A row view starting 8 bytes into the buffer keeps a legal row stride but not a legal bulk-copy source.""" + x, w = _build(2, 128, 6144, 3) + flat = torch.empty(2 * 6144 + 4, dtype=torch.bfloat16, device="cuda") + shifted = flat[4:].view(2, 6144) + shifted.copy_(x) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(tc): + tc(shifted, w, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..59201301c623 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/m3_embedding_firstnorm entry: the decode embedding gather + the first layer's RMSNorm. + +The table is the model's full bf16 [200064, 6144] embedding shape (random values at the checkpoint's scale), eps the +checkpoint's 1e-6. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + m3_embedding_firstnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" + +H = 6144 +EPS = 1e-6 + + +@functools.lru_cache(maxsize=None) +def _operands(): + g = torch.Generator(device="cuda").manual_seed(20261013) + table = (0.02 * torch.randn(VOCAB, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return table, weight + + +def _buffers(rows): + return (torch.empty(rows, H, dtype=torch.bfloat16, device="cuda"), + torch.empty(rows, H, dtype=torch.bfloat16, device="cuda")) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + rows, id_dtype = cell.spec["rows"], cell.spec["id_dtype"] + table, weight = _operands() + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device="cuda").to(id_dtype) + raw, out = _buffers(rows) + with validating(op): + normed, residual = op(ids, table, weight, EPS, raw, out) + assert normed.data_ptr() == out.data_ptr() and residual.data_ptr() == raw.data_ptr() + ref_normed, ref_raw = op.reference(ids, table, weight, EPS, raw, out) + # The gather is a copy: exact. + assert torch.equal(raw.view(torch.int16), ref_raw.view(torch.int16)) + # The norm: the kernel's fp32 sum of squares (warp butterflies, then across 4 warps) and fast rsqrt against an + # fp64 reference, then x * rstd * w rounded once to bf16. The fp32 side is ~1e-6 relative off, so the two roundings + # disagree only next to a bf16 rounding boundary: one ulp at most per element, a small fraction in aggregate. + # Measured over these cells (3 seeds each, GB300): at most 0.51 ulp element, 0.026 ulp RMS. + assert_within_ulp(out, ref_normed, element_ulp=1.0, rms_ulp=0.1) + + +def test_ids_outside_the_table_give_zero_rows() -> None: + """Documented, not refused: an id outside [0, V) gathers a zero row whose norm is zero.""" + table, weight = _operands() + ids = torch.tensor([-1, 5, VOCAB], dtype=torch.int64, device="cuda") + raw, out = _buffers(3) + op(ids, table, weight, EPS, raw, out) + assert not raw[0].any() and not raw[2].any() and not out[0].any() and not out[2].any() + assert torch.equal(raw[1], table[5]) + + +@pytest.mark.parametrize("rows", range(1, 17)) +def test_against_the_stock_embedding_and_rmsnorm(rows) -> None: + """What this target certifies on SM107 against the built-in path it replaces (the model's Embedding gather, then + the first layer's RMSNorm module), including a constant row, large and tiny rows and both vocab ends: the gather + bit for bit, the norm within one bf16 ulp (the sum-of-squares tree differs from flashinfer's: measured ~5e-6 of + the elements one ulp off over 12000 random rows).""" + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + table, weight = _operands() + table = table.clone() + g = torch.Generator(device="cuda").manual_seed(rows) + table[1] = 3.0 + table[2] = (300.0 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + table[3] = (1e-4 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + norm = RMSNorm(hidden_size=H, eps=EPS, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + norm.weight.copy_(weight) + for ids in ( + torch.randint(0, VOCAB, (rows,), generator=g, device="cuda"), + # The special rows (constant, large, tiny, both vocab ends), repeated to the row count past 8. + torch.tensor(([1, 2, 3, 0, VOCAB - 1, 1, 2, 3] * 4)[:rows], device="cuda"), + ): + raw, out = _buffers(rows) + normed, residual = op(ids, table, weight, EPS, raw, out) + x = torch.nn.functional.embedding(ids, table) + assert torch.equal(residual.view(torch.int16), x.view(torch.int16)) + stock = norm(x) + off = normed != stock + assert int(off.sum()) <= 4, f"{int(off.sum())} normed elements differ from stock" + assert_within_ulp(normed, stock, element_ulp=1.0, rms_ulp=0.05) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..158aa26fc6e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs shared by the tp4 target's spec entry tests: TP4 dimensions, single-device stand-ins for the cross-rank +exchanges, and the comparison bands.""" + +from __future__ import annotations + +import zlib + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) + +#: MiniMax-M3 at TP4: hidden size, vocabulary, the target head's shard and RMSNorm epsilon. +HIDDEN = 6144 +VOCAB = 200064 +WORLD = 2 +TARGET_SHARD = VOCAB // WORLD +EPS = 1e-6 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): + """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" + uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") + keys = torch.zeros(slot_rows, dtype=torch.int64, device="cuda") + counters = torch.tensor([0, call], dtype=torch.int32, device="cuda") + return uc, keys, counters + + +def plant_peers( + uc: torch.Tensor, + counters: torch.Tensor, + own_best: torch.Tensor, + rank: int, + world: int, + shard: int, + g: torch.Generator, + slot_rows: int = SHARD_ROWS, + ties: bool = False, +) -> None: + """Write the other ranks' keys for each row into this rank's slots, as their multicast pushes would. + + Each peer's maximum is the own row maximum shifted by +-0.25 (well clear of the accumulation-order band), so + some rows are won by a peer and some by this rank; with ``ties`` a peer on an even row reports exactly the own + maximum, and the lower global index must win. + """ + if world == 1: + return + m = own_best.shape[0] + half = int(counters[1].item()) & 1 + for r in range(world): + if r == rank: + continue + sign = torch.where(torch.rand(m, generator=g, device="cuda") < 0.5, -1.0, 1.0) + value = own_best.float() + 0.25 * sign + if ties: + value = torch.where(torch.arange(m, device="cuda") % 2 == 0, own_best.float(), value) + index = r * shard + torch.randint(0, shard, (m,), generator=g, device="cuda") + keys = make_keys(value, index) + for t in range(m): + uc[(half * slot_rows + t) * world + r] = keys[t] + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..15758e2645e2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_capture_norm entry (Eagle3 capture with fc_norm folded in).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_capture_norm import ( # noqa: E501 + m3_capture_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" + +LAYERS = 3 + + +def _case(spec, seed): + g = gen(seed) + m, slot = spec["m"], spec["slot"] + h = randn((m, HIDDEN), g, 2.0) + r = randn((m, HIDDEN), g, 2.0) if spec["residual"] else None + w = (1.0 + randn((HIDDEN,), g, 0.1).float()).to(torch.bfloat16) + buf = torch.full((m, LAYERS * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + dst = buf[:, slot * HIDDEN : (slot + 1) * HIDDEN] + return (h, r, w, EPS, dst), buf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer + untouched.""" + args, buf = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + slot = cell.spec["slot"] + others = torch.cat([buf[:, : slot * HIDDEN], buf[:, (slot + 1) * HIDDEN :]], dim=1) + assert bool((others == 7.0).all()) + + +def test_guard_refuses_a_strided_weight() -> None: + (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + w2 = torch.stack([w, w], dim=1)[:, 0] + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(h, r, w2, eps, dst) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py new file mode 100644 index 000000000000..fde15dd95781 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The lazily created exchange workspaces refuse to be created inside a CUDA-graph capture (K3 audit #150, check #10). + +The draft-step and sampled-accept workspaces are created collectively (both ranks, a multicast buffer) on the first +eager call. Creating one inside a capture would allocate and synchronize mid-capture, so a capture that finds it missing +reports a miss (``kernels/m3_report.capture_missed``) and gets None (the caller then takes its stock path) instead. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( + m3_draft_step, + m3_report, + m3_sampled_accept, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU" +) + + +def _refuse(*args, **kwargs): + raise AssertionError("a workspace must not be created inside a CUDA-graph capture") + + +@pytest.mark.parametrize( + "module, creator", + [ + (m3_draft_step, "create_workspace"), + (m3_sampled_accept, "SampledAcceptWorkspace.create"), + ], + ids=["draft_step", "sampled_accept"], +) +def test_a_capture_without_the_workspace_misses_instead_of_creating_it( + monkeypatch, module, creator +): + monkeypatch.setattr(module, "_WS", None) + owner, _, attr = creator.rpartition(".") + monkeypatch.setattr(getattr(module, owner) if owner else module, attr, _refuse) + misses = [] + monkeypatch.setattr(m3_report, "capture_missed", lambda *args: misses.append(args)) + + device = torch.device("cuda", torch.cuda.current_device()) + found = [] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + found.append(module.workspace(None, device)) + torch.cuda.current_stream().wait_stream(stream) + + assert found == [None] + assert module._WS is None + assert len(misses) == 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..b255f786fc5e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_chunk_norm entry (the Eagle3 fc_norm over the captured layers).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_chunk_norm import ( # noqa: E501 + m3_chunk_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, c, h = spec["m"], spec["c"], spec["h"] + buf = randn((m, c * h + spec["extra_cols"]), g, 2.0) + src = buf[:, : c * h + spec["extra_cols"]] + w = (1.0 + randn((c, h), g, 0.1).float()).to(torch.bfloat16) + dst = torch.full((m, c * h), 7.0, dtype=torch.bfloat16, device="cuda") + return (src, w, EPS, dst, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + + +def test_guard_refuses_a_strided_weight() -> None: + src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() + wt.copy_(w) + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(src, wt, eps, dst, pdl) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py new file mode 100644 index 000000000000..30931023b5ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the drafter boundary push (modeling.DRAFT_AR_PUSH / DRAFT_AR_PUSH_DIRECT) on one rank: the push variants +of the split-K copy (kernels/m3_lowm_splitk) and the direct copy (kernels/m3_lowm_direct) write FlashInfer's output +bit for bit, and store exactly the boundary all-reduce's canonical exchange words of that output: word (t, column pair) += bf16 lo | bf16 hi << 16 with -0.0 halves as +0.0, at parity * maxt * world * H / 2 + (t * world + rank) * H / 2 + +column / 2, the parity being the (t, column / (H / 8)) flag's low bit; every other word stays untouched. Eager and in a +CUDA graph, both parities, at each certified push cell of the TP2 table (the 1-row direct and 4-row split-K o_proj / +down).""" + +import pytest +import torch +from flashinfer.gemm.kernels.dense_bf16_gemm_direct import DirectTactic as FiDirect +from flashinfer.gemm.kernels.dense_bf16_gemm_direct import run_direct_dense as fi_direct +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiSplitK +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_splitk + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_lowm_direct, + m3_lowm_splitk, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + tactic_for, +) + +assert torch.cuda.is_available(), "the push variants require a CUDA device" + +H, MAXT, EMPTY = 6144, 8, -2147483648 +# The boundary linears (output width H): o_proj and down, each certified row count on split-K or an even direct width, +# up to the 8-token bucket the push arms on (kernels/m3_draft_lowm.arm_push); the batch-4/8 rows past it run unpushed. +CELLS = [ + (nk[1], m, *tactic_for(nk[0], nk[1], m)) + for nk, rows in sorted(ROWS.items()) + if nk[0] == H and nk[1] != 18432 # o_proj / down, not the fc + for m in rows + if m <= MAXT + and (tactic_for(nk[0], nk[1], m)[0] == "s" or tactic_for(nk[0], nk[1], m)[1][1] % 2 == 0) +] + + +def _expected_words(y: torch.Tensor, flags: torch.Tensor, world: int, rank: int) -> torch.Tensor: + m = y.shape[0] + half = MAXT * world * (H // 2) + ws = torch.full((2 * half,), EMPTY, dtype=torch.int32, device="cuda") + h16 = y.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, H // 2) + for t in range(m): + for c in range(8): + par = int(flags[t * 8 + c].item()) & 1 + lo = par * half + (t * world + rank) * (H // 2) + c * (H // 16) + ws[lo : lo + H // 16] = words[t, c * (H // 16) : (c + 1) * (H // 16)] + return ws + + +@pytest.mark.parametrize("parity", [0, 1]) +@pytest.mark.parametrize("cell", CELLS, ids=[f"{k}@{m}{kind}" for k, m, kind, _ in CELLS]) +def test_push_words_and_output(cell, parity) -> None: + k, m, kind, tactic = cell + world, rank = 2, 1 + g = torch.Generator(device="cuda").manual_seed(k + m + parity) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + x[0, :8] = ( + 0.0 # zero / signed-zero outputs exercise the -0 canonicalization through the weights below + ) + w = (torch.randn(H, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + w[:4] = 0.0 + w[4:8] = -0.0 + flags = ( + torch.randint(0, 7, (MAXT * 8,), generator=g, device="cuda", dtype=torch.int32) * 2 + parity + ) + ref = torch.empty(m, H, dtype=torch.bfloat16, device="cuda") + if kind == "d": + fi_direct(x, w.t(), ref, True, FiDirect(*tactic)) + else: + fi_splitk(x, w.t(), None, ref, True, FiSplitK(*tactic)) + + def run(out, ws): + ar = (ws.data_ptr(), flags.data_ptr(), rank, world, MAXT) + if kind == "d": + m3_lowm_direct.run_direct_dense_push( + x, w.t(), out, True, m3_lowm_direct.DirectTactic(*tactic), ar + ) + else: + m3_lowm_splitk.run_splitk_dense_push( + x, w.t(), out, True, m3_lowm_splitk.SplitKTactic(*tactic), ar + ) + + want = _expected_words(ref, flags, world, rank) + for graph in (False, True): + out = torch.empty(m, H, dtype=torch.bfloat16, device="cuda") + ws = torch.full_like(want, EMPTY) + if graph: + run(out, ws) # compile outside capture + torch.cuda.synchronize() + ws.fill_(EMPTY) + gr = torch.cuda.CUDAGraph() + with torch.cuda.graph(gr): + run(out, ws) + ws.fill_(EMPTY) + gr.replay() + else: + run(out, ws) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), f"output (graph={graph})" + assert torch.equal(ws, want), f"exchange words (graph={graph})" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..9669fc829def --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_input entry (embedding + both RMSNorms + concat in one launch), against its reference and +against the in-tree drafter's modules.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_input import ( # noqa: E501 + m3_draft_input as op, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + VOCAB, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, h = spec["m"], spec["h"] + table = randn((VOCAB, h), g) + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + if spec["invalid_id"]: + ids[0] = -1 + hidden = randn((4 * m if spec["gather"] else m, h), g) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 2 if spec["gather"] else None + w_e = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + w_h = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + out = torch.full((m, 2 * h), 7.0, dtype=torch.bfloat16, device="cuda") + residual = torch.full((m, h), 7.0, dtype=torch.bfloat16, device="cuda") + return (ids, table, hidden, rows, w_e, w_h, EPS, EPS, out, residual, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Normalized halves within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic in a + different summation order, rsqrtf); the gathered residual and the zero embedding of an invalid id exact.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out, residual = op(*args) + ref_out, ref_res = op.reference(*args) + torch.cuda.synchronize() + assert_bf16_close(out, ref_out) + assert torch.equal(residual, ref_res) + if cell.spec["invalid_id"]: + assert not bool(out[0, : cell.spec["h"]].float().any()) + + +def test_guard_refuses_rows_past_hidden() -> None: + gather = next(c for c in op.CELLS if c.spec["gather"]) + args = list(_case(gather.spec, 1)) + args[3] = args[3] + 100 + with pytest.raises(AssertionError, match="rows must index hidden"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("m,gather", [(4, False), (32, False), (8, True)], ids=str) +def test_the_in_tree_draft_input(m, gather) -> None: + """Against the in-tree drafter's own modules: the Embedding lookup, then the layer's input_layernorm and + hidden_norm (in-tree RMSNorm) and the concat of Eagle3DecoderLayer.forward. Within one bf16 ulp on at most 0.1% of + the elements (the RMS statistic summed in another order); the residual (the carried hidden rows) exact.""" + g = gen(seed_of(f"in-tree draft input {m} {gather}")) + h = HIDDEN + embed = Embedding(VOCAB, h, dtype=torch.bfloat16).cuda() + embed.weight.data = randn((VOCAB, h), g) + norms = [] + for _ in range(2): + norm = RMSNorm(hidden_size=h, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + norms.append(norm) + input_layernorm, hidden_norm = norms + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + hidden = randn((4 * m if gather else m, h), g, 2.0) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 3 if gather else None + carried = hidden[rows] if gather else hidden + stock = torch.cat([input_layernorm(embed(ids)), hidden_norm(carried)], dim=-1) + out, residual = m3_spec_tail.draft_input( + ids, embed.weight, hidden, rows, input_layernorm.weight, EPS, hidden_norm.weight, EPS + ) + torch.cuda.synchronize() + assert_bf16_close(out, stock) + assert torch.equal(residual, carried) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_push_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_push_take.py new file mode 100644 index 000000000000..523385b5dd06 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_push_take.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter boundary push bookkeeping (``M3Eagle3Draft.arm_push`` / ``take_pushed``, ``DRAFT_AR_PUSH``): the +boundary all-reduce runs prepushed only when the armed call pushed, and only on that call's own output; any other +input raises instead of reducing rows the exchange does not hold. Host state only (no GPU).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) + + +@pytest.fixture +def draft(): + """A drafter with only its push bookkeeping (the module and its weights are not built).""" + d = object.__new__(modeling.M3Eagle3Draft) + d._push_arm, d._pushed = None, 0 + return d + + +def test_not_pushed_is_false_and_disarms(draft) -> None: + draft._push_arm = (1, 2, 3, 0, 8, 8) + assert draft.take_pushed(torch.zeros(4)) is False + assert draft._push_arm is None + + +def test_pushed_output_is_true_once(draft) -> None: + out = torch.zeros(4, 8) + draft._pushed = out.data_ptr() + assert draft.take_pushed(out) is True + assert draft._pushed == 0 + assert draft.take_pushed(out) is False + + +def test_other_input_raises(draft) -> None: + out = torch.zeros(4, 8) + other = out.clone() + draft._pushed = out.data_ptr() + with pytest.raises(RuntimeError, match="not the boundary all-reduce's input"): + draft.take_pushed(other) + assert draft._pushed == 0 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py new file mode 100644 index 000000000000..a817e8ce00dc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_silu entry (the drafter MLP's plain SwiGLU as one C++ launch): bit for bit the Triton +silu_and_mul at its defaults, over every bf16 gate and every bf16 up pattern.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_silu import ( # noqa: E501 + m3_draft_silu as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" + +ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) + + +def _rows(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """[rows, 2d] gate | up from equally long flat gate and up values, d = 9216 (the TP2 drafter shard).""" + d = 9216 + n = gate.numel() + rows = -(-n // d) + pad = rows * d - n + g = torch.cat([gate, gate[:pad]]).view(rows, d) + u = torch.cat([up, up[:pad]]).view(rows, d) + return torch.cat([g, u], dim=1).contiguous() + + +def _assert_bitwise(x: torch.Tensor, what: str) -> None: + got = op(x) + assert got is not None, what + ref = op.reference(x) + torch.cuda.synchronize() + diff = got.view(torch.int16) != ref.view(torch.int16) + assert not bool(diff.any()), f"{what}: {int(diff.sum())} of {diff.numel()} differ" + + +@pytest.mark.parametrize("up", [0.0, 1.0, -1.0, 3.5, -0.0, 1e-3], ids=lambda v: f"up={v}") +def test_every_gate_pattern_is_bitwise_the_triton_kernel(up) -> None: + gate = ALL.cuda() + _assert_bitwise(_rows(gate, torch.full_like(gate, up)), f"gate sweep, up {up}") + + +@pytest.mark.parametrize("gate", [1.0, -1.0, 0.5, -3.0, 20.0, -0.0], ids=lambda v: f"gate={v}") +def test_every_up_pattern_is_bitwise_the_triton_kernel(gate) -> None: + up = ALL.cuda() + _assert_bitwise(_rows(torch.full_like(up, gate), up), f"up sweep, gate {gate}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed) -> None: + g = gen(100 + seed) + n = 1 << 22 + gate = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + up = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + _assert_bitwise(_rows(gate.view(torch.bfloat16), up.view(torch.bfloat16)), f"seed {seed}") + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The drafter's gate_up output range (and past it), in a CUDA graph as the drafter runs it.""" + g = gen(seed_of(cell.why)) + rows, d = cell.spec["rows"], cell.spec["d"] + x = (4.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :6] = torch.tensor([60.0, -60.0, 0.0, -0.0, 88.0, -88.0], dtype=torch.bfloat16) + certifying(op) + ref = op.reference(x) + eager = op(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(x) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(eager.view(torch.int16), ref.view(torch.int16)) + assert torch.equal(captured.view(torch.int16), ref.view(torch.int16)) + + +def test_declines_what_it_does_not_cover() -> None: + x = torch.zeros(2, 2 * 9216, dtype=torch.bfloat16, device="cuda") + assert op(x[:, 8:]) is None # strided rows + assert ( + op(torch.zeros(2, 2 * 12, dtype=torch.bfloat16, device="cuda")) is None + ) # d not a multiple of 8 + assert op(x.float()) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..7d0cc68a5011 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,814 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_step entry (one fused greedy draft pass: the BF16 draft head, streamed from its lossless +12-bit coding or as BF16 fragments, the cross-rank argmax and the draft loop's bookkeeping), and its tokens against the +in-tree drafter's head and greedy draft sampling. + +The in-tree path: each rank's LM head shard GEMM (BF16 out), ``LogitsProcessor``'s float cast, then either the argmax +of the gathered logits (``torch.argmax``) or ``greedy_sample_draft_with_tp_gather``'s per-rank ``torch.max`` and the +pick among the ranks' (index, value) pairs. The edge cases use integer data, whose fp32 sums are exact in any order, +so the kernel and the in-tree path see the same BF16 logits and must agree bit for bit; every in-tree comparison runs +with the coded head (the default) and with the BF16 fragments, and on integer data the coded head's winning keys are +the BF16 fragments' bit for bit. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.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_107__tp2.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_LEN, + DRAFT_SHARD, + HIDDEN, + MAX_TOKENS, + bf16_fragments, + from_fragments, + head_weight, + shard_logits, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_step import ( # noqa: E501 + m3_draft_step as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + WORLD, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" + +T = DRAFT_LEN + 1 +_LOW32 = 0xFFFFFFFF + +# ------------------------------------------------------------------ the per-row rule against the in-tree token +# +# The kernel and the in-tree head (a cuBLAS BF16 GEMM) both round an FP32 sum of the same exact BF16 products to BF16, +# in different summation orders. A row's token can differ from the in-tree one only where some candidate's FP32 sum +# can land on either side of a BF16 rounding midpoint. Each row is held to: +# +# kernel token == in-tree token, or the row is provably undecidable: with the fp64 reference logit L_i (exact BF16 +# products, fp64 sum) and the FP32-order bound E_i = gamma_K * sum_k |w_ik * x_k|, the BF16 value of any FP32 sum of +# candidate i lies in [lo_i, hi_i] = [bf16(L_i - E_i), bf16(L_i + E_i)]; the kernel's token must win for some +# choice of values within those intervals (larger value, then the lower global id), and a BF16 rounding midpoint +# must lie within E of the kernel's or the in-tree token's L. +# +# gamma_K = K * 2^-22 is la-xh2's bound (4x the round-to-nearest unit roundoff 2^-24 per addition: any summation +# order, and a truncating tensor-core accumulator); K counts the row's non-zero products (adding an exact zero does +# not round). lo / hi take one more FP32 step outward before the BF16 rounding, which covers rounding L -+ E to FP32. +# Each undecidable row is reported as a warning with its midpoint distances, and each test reports its row count; any +# other mismatch fails. +GAMMA_UNIT = 2.0**-22 + + +def _bf16_outward(v: torch.Tensor, direction: float) -> torch.Tensor: + f = v.float() + f = torch.nextafter(f, torch.full_like(f, direction)) + b = f.to(torch.bfloat16).float() + return torch.where(b == 0, torch.zeros_like(b), b) + + +def bounds(h: torch.Tensor, w: torch.Tensor): + """(fp64 logits L, bound E, lo, hi) [rows, N] of the head input rows h against the head rows w.""" + cols = (h != 0).any(dim=0) # zero columns add exact zeros + hd, wd = h[:, cols].double(), w[:, cols].double() + logits = hd @ wd.t() + bound = (h != 0).sum(dim=1, keepdim=True).double() * GAMMA_UNIT * (hd.abs() @ wd.abs().t()) + del wd + inf = float("inf") + return logits, bound, _bf16_outward(logits - bound, -inf), _bf16_outward(logits + bound, inf) + + +def midpoint_distance(v: float) -> float: + """Distance of v to the nearest BF16 rounding midpoint.""" + b = torch.tensor(v, dtype=torch.float64).float().to(torch.bfloat16).float() + if float(b) == 0.0: + return abs(v) + ulp = 2.0 ** (int(torch.frexp(b).exponent) - 8) + return min(abs(v - (float(b) + ulp / 2)), abs(v - (float(b) - ulp / 2))) + + +def _winnable(tok: int, ids: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Whether candidate ``tok`` wins the row for some values within the candidates' [lo, hi].""" + k = int((ids == tok).nonzero()[0, 0]) + others = ids != tok + beat = (hi[k] > lo[others]) | ((hi[k] == lo[others]) & (tok < ids[others])) + return bool(beat.all()) + + +def check_row(name: str, kernel_tok: int, stock_tok: int, ids, logits, bound, lo, hi) -> bool: + """The per-row rule for one row over the candidates ``ids`` [C] (logits / bound / lo / hi [C]); True when the + row is undecidable (reported as a warning), False when the tokens match; a mismatch outside the bound fails.""" + if kernel_tok == stock_tok: + return False + assert _winnable(kernel_tok, ids, lo, hi), ( + f"{name}: token {kernel_tok} differs from the in-tree {stock_tok} outside the FP32-order bound" + ) + ik = int((ids == kernel_tok).nonzero()[0, 0]) + i_stock = int((ids == stock_tok).nonzero()[0, 0]) + lk, ls, ek, es = ( + float(logits[ik]), + float(logits[i_stock]), + float(bound[ik]), + float(bound[i_stock]), + ) + dk, ds = midpoint_distance(lk), midpoint_distance(ls) + assert dk <= ek or ds <= es, f"{name}: neither logit is within its bound of a BF16 midpoint" + warnings.warn( + f"[draft-step undecidable] {name}: kernel {kernel_tok} (L {lk:.9g}, midpoint distance {dk:.3e}, bound " + f"{ek:.3e}) vs in-tree {stock_tok} (L {ls:.9g}, midpoint distance {ds:.3e}, bound {es:.3e})" + ) + return True + + +def report(name: str, rows: int, undecidable: int) -> None: + warnings.warn(f"[draft-step rows] {name}: {rows} rows, {undecidable} undecidable") + + +def key_value(keys: torch.Tensor) -> torch.Tensor: + """The fp32 value a key orders by (the inverse of the kernel's order-preserving image).""" + u = (keys >> 32) & _LOW32 + bits = torch.where(u >= 0x80000000, u - 0x80000000, _LOW32 - u) + bits = torch.where(bits >= 2**31, bits - 2**32, bits) # as a signed 32-bit pattern + return bits.to(torch.int32).view(torch.float32) + + +def key_index(keys: torch.Tensor) -> torch.Tensor: + return _LOW32 - (keys & _LOW32) + + +def _bookkeeping(nreq: int, step: int, n_x_rows: int, g: torch.Generator): + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if step == 0 + else None + ) + seq_lens = torch.full((nreq,), T, dtype=torch.int32, device="cuda") if step == 0 else None + next_new = torch.full((nreq, T), -5, dtype=torch.int32, device="cuda") + pos_src = torch.randint(100, 3000, (n_x_rows,), generator=g, device="cuda", dtype=torch.int32) + kv_lens = torch.randint(200, 3000, (nreq,), generator=g, device="cuda", dtype=torch.int32) + ids_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + pos_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + return num_acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next + + +def _head(w: torch.Tensor, coded: bool): + """(wf, (exc_off, exc_pos, exc_w, exp_base)) of a BF16 shard: its 12-bit coding (checked to round-trip) or its + BF16 fragments.""" + if not coded: + return bf16_fragments(w), (None, None, None, 0) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + return packed, (exc_off, exc_pos, exc_w.view(-1), base) + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=(None, None, None, 0), +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + coded head's (exc_off, exc_pos, exc_w, exp_base).""" + acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next = _bookkeeping( + nreq, step, x.shape[0], g + ) + if num_acc is None: + num_acc = acc + if uc is None: + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + out_keys = torch.zeros(nreq, dtype=torch.int64, device="cuda") + tiles = wf.shape[0] if coded[0] is None else coded[0].numel() - 1 + return (x, x_rows, nreq, wf, rank * tiles * 16, keys, counters, uc, uc.data_ptr(), rank, world, + next_new[:, 1:], step, ids_next, pos_src, pos_next, kv_lens, num_acc, seq_lens, DRAFT_LEN, num_sms(), + True, out_keys, *coded) # fmt: skip + + +def _case(spec, seed): + g = gen(seed) + nreq, step, rank, world, vs, k = ( + spec["nreq"], + spec["step"], + spec["rank"], + spec["world"], + spec["vs"], + spec["k"], + ) + pass0 = step == 0 + x = randn((nreq * T if pass0 else nreq, k), g) + wf, coded = _head(randn((vs, k), g, 0.02), spec["coded"]) + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if pass0 + else None + ) + x_rows = (torch.arange(nreq, device="cuda") * T + num_acc.long() - 1) if pass0 else None + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + h = x[x_rows] if pass0 else x + own = shard_logits(h, wf, *coded).max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=MAX_TOKENS, ties=spec.get("ties", False) + ) + return _args(x, x_rows, nreq, wf, rank, world, step, g, uc, keys, counters, num_acc, coded) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The bookkeeping is exact on every row (next positions, KV and sequence lengths). Each row's token is the + in-tree one (this rank's BF16 shard logits from the head's GEMM, their first maximum, against the peers' keys) or + the row is provably undecidable (check_row); the winning key's index is the token, and its value lies in the + token's [lo, hi] (or is the peer's pushed key when a peer wins).""" + args = _case(cell.spec, seed_of(cell.why)) + x, x_rows, nreq, wf, offset, _, counters, uc = args[:8] + rank, world = args[9], args[10] + half = int(counters[1].item()) & 1 + pushed = ( + uc[half * MAX_TOKENS * world : (half * MAX_TOKENS + nreq) * world].view(nreq, world).clone() + ) + h = x[x_rows] if x_rows is not None else x[:nreq] + w = head_weight(wf, x.shape[1], *args[23:27]) + best, idx = local_argmax(torch.nn.functional.linear(h, w).float(), offset) + stock = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, MAX_TOKENS) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert args[6].tolist() == [0, 1] and not bool(args[5].any()), "counters / keys not re-armed" + step = args[12] + others = [c for c in range(args[11].shape[1]) if c != step] + assert torch.equal(args[11][:, others], ref["draft_out"][:, others]) + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + assert torch.equal(args[15], ref["pos_next"]) + assert torch.equal(args[16], ref["kv_lens"]) + if args[18] is not None: + assert torch.equal(args[18], ref["seq_lens"]) + # Candidates per row: this rank's head rows (fp64 bounds) and each peer's pushed key (an exact value). + logits, bound, lo, hi = bounds(h, w) + own_ids = offset + torch.arange(w.shape[0], device="cuda") + peers = [p for p in range(world) if p != rank] + undecidable = 0 + for r in range(nreq): + pv = key_value(pushed[r, peers]).double() + ids = torch.cat([own_ids, key_index(pushed[r, peers])]) + undecidable += check_row( + f"{cell.why[:40]} row {r}", + int(args[13][r]), + int(stock[r]), + ids, + torch.cat([logits[r], pv]), + torch.cat([bound[r], torch.zeros_like(pv)]), + torch.cat([lo[r], pv.float()]), + torch.cat([hi[r], pv.float()]), + ) + report(cell.why[:40], nreq, undecidable) + out_keys = args[22] + assert torch.equal(key_index(out_keys), args[13].long()), "the winning key's index is the token" + got = key_value(out_keys) + for r in range(nreq): + c = int(args[13][r]) - offset + if 0 <= c < w.shape[0]: + assert float(lo[r, c]) <= float(got[r]) <= float(hi[r, c]), ( + f"row {r}: winning value outside [lo, hi]" + ) + else: + assert int(out_keys[r]) in pushed[r].tolist(), ( + f"row {r}: a peer won with a key it did not push" + ) + + +def test_fragments_round_trip() -> None: + g = gen(seed_of("fragments")) + w = randn((64, 96), g) + wf = bf16_fragments(w) + assert wf.shape == (4, 6, 256) + assert torch.equal(from_fragments(wf), w) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 3) # K 96: no coding, the BF16 fragments + assert not prepared.coded + assert torch.equal(prepared.wf, wf) and prepared.vocab_offset == 3 * 64 and prepared.weight is w + + +# ------------------------------------------------------------------ against the in-tree drafter, bit for bit + + +def _stock_tokens(shards: list, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The in-tree greedy draft tokens of the rows x over the vocab-sharded BF16 head: (argmax of the gathered float + logits, greedy_sample_draft_with_tp_gather's per-rank (index, value) pick).""" + logits = [torch.nn.functional.linear(x, w).float() for w in shards] + gathered = torch.argmax(torch.cat(logits, dim=-1), dim=-1).to(torch.int32) + combined = [] + for rank, shard_logit in enumerate(logits): + worker = SimpleNamespace(mapping=SimpleNamespace(tp_rank=rank)) + combined.append(Eagle3OneModelWorker._get_local_max_and_combined(worker, shard_logit)) + pick = Eagle3OneModelWorker._get_draft_tokens_from_gathered + pick = getattr(pick, "_torchdynamo_orig_callable", pick) + tp_gather = pick(None, torch.cat(combined, dim=-1)).to(torch.int32) + return gathered, tp_gather + + +def _fused_tokens(shards: list, x: torch.Tensor, coded: bool, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's four-rank pass with + its peers' keys in its slots, as their multicast pushes would put them.""" + g = gen(seed_of("fused tokens")) + nreq = x.shape[0] + heads = [_head(w, coded) for w in shards] + own_keys = [] + for rank, (wf, cd) in enumerate(heads): + args = list(_args(x, None, nreq, wf, rank, 1, step, g, coded=cd)) + args[9] = 0 # one-rank run: this rank's slot only, global ids of `rank` + op(*args) + own_keys.append(args[22].clone()) + tokens = [] + for rank, (wf, cd) in enumerate(heads): + uc, keys, counters = head_workspace(WORLD, slot_rows=MAX_TOKENS) + for peer in range(WORLD): + if peer != rank: + for t in range(nreq): + uc[t * WORLD + peer] = own_keys[peer][t] + args = _args(x, None, nreq, wf, rank, WORLD, step, g, uc, keys, counters, coded=cd) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + tokens.append(args[13].clone()) + return tokens + + +def _integer_head(g: torch.Generator, rows: int = WORLD * DRAFT_SHARD) -> torch.Tensor: + """A [rows, 6144] head of small integers over the first 4 columns (zero elsewhere): its logits against integer + rows are exact integers in any summation order, many of which share a BF16 value.""" + w = torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, :4] = torch.randint(-60, 61, (rows, 4), generator=g, device="cuda").to(torch.bfloat16) + return w + + +def _shards(w: torch.Tensor) -> list: + return [w[r * DRAFT_SHARD : (r + 1) * DRAFT_SHARD].contiguous() for r in range(WORLD)] + + +def _assert_stock(w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None) -> None: + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + if expect is not None: + assert gathered.tolist() == expect, (gathered.tolist(), expect) + for rank, tokens in enumerate(_fused_tokens(shards, x, coded)): + assert torch.equal(tokens, gathered), ( + f"rank {rank}: {tokens.tolist()} vs stock {gathered.tolist()}" + ) + + +def _ones_rows(nreq: int) -> torch.Tensor: + x = torch.zeros((nreq, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[:, :4] = 1.0 + return x + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_forced_bf16_ties_take_the_lowest_global_id(coded) -> None: + """A head row copied onto every rank that wins every row of ones (logit 244, every other at most 240): the lowest + global id must win, inside a shard and across ranks. Then integer rows in [-2, 2], whose logits (up to 480 in + magnitude, BF16 spacing 2 above 256) collide often, exactly or after BF16 rounding.""" + g = gen(seed_of("forced ties")) + w = _integer_head(g) + win = torch.tensor([61.0, 61.0, 61.0, 61.0], device="cuda").to(torch.bfloat16) + for gid in (DRAFT_SHARD + 5000, DRAFT_SHARD + 77, DRAFT_SHARD + 9000, 95000): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), coded, expect=[95000] * 5) + # Random integer rows: ties after BF16 rounding happen naturally. + x = torch.randint(-2, 3, (8, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + x[:, 4:] = 0 + _assert_stock(w, x, coded) + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_signed_zero_maxima_take_the_lowest_global_id(coded) -> None: + """Rows whose maximum is zero: an all-zero head input (every logit +-0) and logits <= 0 with zeros by cancellation + on several ranks. -0 and +0 compare equal, so the lowest global id reaching zero wins.""" + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 0] = -1.0 + w[:, 1] = -1.0 + for gid in (DRAFT_SHARD + 5, DRAFT_SHARD + 40000, DRAFT_SHARD + 70007, 49999): + w[gid, 0] = 1.0 + w[gid, 1] = -1.0 + x = torch.zeros((3, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[1, :2] = 1.0 + x[2, :2] = 2.0 + x[2, 2] = -1.0 # w[:, 2] is zero: a -0 product on every row + _assert_stock(w, x, coded, expect=[0, 49999, 49999]) + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + (1000, 1001, "lower"), # 1001 rounds to 1000 in BF16: a tie, the lower id wins + (1001, 1002, "lower"), # 1001 and the tie-to-even midpoint 1002 both round to 1000 + (1000, 1004, "higher"), # adjacent BF16 values: the larger wins + (1003, 1000, "lower"), # 1003 rounds up to 1004 > 1000 + ], + ids=["one_fp32_step", "midpoint", "adjacent_bf16", "rounds_up"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where, coded) -> None: + """Top-2 logits one to four fp32 steps apart: the winner is the one the BF16 logits decide (the in-tree logits' + dtype), not the fp32 sum. The pair sits inside rank 1's shard or on ranks 0 and 1; every other logit is at most + 240 in magnitude.""" + g = gen(seed_of(f"near top2 {lower} {higher} {where}")) + w = _integer_head(g) + if where == "in_shard": + lo_id, hi_id = DRAFT_SHARD + 100, DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD - 1000, DRAFT_SHARD + 3 + for gid, value in ((lo_id, lower), (hi_id, higher)): + q, r = divmod(value, 4) # four integers <= 251, exact in BF16 + w[gid, :4] = torch.tensor([q + (i < r) for i in range(4)], device="cuda").to(torch.bfloat16) + expect = lo_id if winner == "lower" else hi_id + _assert_stock(w, _ones_rows(2), coded, expect=[expect, expect]) + + +# ------------------------------------------------------------------ non-integer rows at a BF16 rounding midpoint + +#: A BF16 rounding midpoint (between 6.0 and 6.03125; a tie there rounds to the even 6.0) and the FP32 spacing there. +MIDPOINT = 6.0 + 2.0**-6 +FP32_ULP = 2.0**-21 + + +def _near_midpoint_head(g: torch.Generator, targets: dict): + """A full-vocabulary head [2 * 100032, 6144] and one head-input row x with 65 non-zero entries: 1, 2^-12, 2^-20 + and 2^-28 on columns 0..3 and 61 random non-integer BF16 values. Every head row carries random non-integer BF16 + values on those 61 columns (logits about 0.2 in magnitude); each row of ``targets`` (global id -> fp64 value) also + gets columns 0..3 set so its fp64 logit is the target to within 2^-30 (a coarse value, then three finer + corrections, each BF16-rounded).""" + x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device="cuda") + scales = (1.0, 2.0**-12, 2.0**-20, 2.0**-28) + for col, scale in enumerate(scales): + x[0, col] = scale + x[0, 4:65] = randn((61,), g, 0.5) + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 4:65] = randn((WORLD * DRAFT_SHARD, 61), g, 0.05) + xd = x[0].double() + for gid, target in targets.items(): + rest = float(target - (w[gid].double() * xd).sum()) + for col, scale in enumerate(scales): + v = torch.tensor(rest / scale, dtype=torch.float64).to(torch.bfloat16) + w[gid, col] = v + rest -= float(v) * scale + assert abs(rest) <= 2.0**-30, (gid, rest) + return w, x + + +def _check_fused_against_stock( + name: str, w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None +) -> int: + """Each rank's fused tokens against the in-tree tokens under the per-row rule (full-vocabulary candidates); + ``expect`` (a token per row) additionally pins rows the bound decides. Returns the undecidable rows.""" + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + fused = _fused_tokens(shards, x, coded) + for tokens in fused[1:]: + assert torch.equal(tokens, fused[0]), "the ranks disagree on the global token" + logits, bound, lo, hi = bounds(x, w) + ids = torch.arange(w.shape[0], device="cuda") + undecidable = 0 + for r in range(x.shape[0]): + undecidable += check_row( + f"{name} row {r}", + int(fused[0][r]), + int(gathered[r]), + ids, + logits[r], + bound[r], + lo[r], + hi[r], + ) + if expect is not None: + assert int(gathered[r]) == expect[r] and int(fused[0][r]) == expect[r], ( + name, + int(gathered[r]), + int(fused[0][r]), + expect[r], + ) + others = ids != expect[r] + assert _winnable(expect[r], ids, lo[r], hi[r]), ( + f"{name}: the bound does not let {expect[r]} win" + ) + rivals = [int(i) for i in ids[others][hi[r][others] >= lo[r][expect[r]]]] + assert not any(_winnable(i, ids, lo[r], hi[r]) for i in rivals), ( + f"{name}: the row is not decided" + ) + report(name, x.shape[0], undecidable) + return undecidable + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + ( + -2, + 2, + None, + ), # straddling the midpoint by 2 FP32 ulps: either side is in-tree up to the order + (2, -2, None), + (0, -2, None), # exactly on the midpoint (ties to the even 6.0) against just below it + ( + -2048, + 2048, + "higher", + ), # ~1e-3 either side, well outside the bound: decided by BF16 rounding + ( + -4096, + -2048, + "lower", + ), # both below the midpoint, ~1e-3 apart: a BF16 tie the lower id wins + ], + ids=["straddle", "straddle_reversed", "on_midpoint", "decided_up", "decided_tie"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where, coded) -> None: + """Non-integer rows whose top-2 fp64 logits sit the given number of FP32 ulps from a BF16 rounding midpoint, inside + rank 1's shard or on ranks 0 and 1, every other logit far below. The per-row rule holds on every row; outside the + bound the token is the one the BF16 rounding decides, bit for bit.""" + g = gen(seed_of(f"midpoint {lower} {higher} {where}")) + if where == "in_shard": + lo_id, hi_id = DRAFT_SHARD + 100, DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD - 1000, DRAFT_SHARD + 3 + w, x = _near_midpoint_head( + g, {lo_id: MIDPOINT + lower * FP32_ULP, hi_id: MIDPOINT + higher * FP32_ULP} + ) + expect = None if winner is None else [lo_id if winner == "lower" else hi_id] + _check_fused_against_stock(f"midpoint {where} {lower:+d}/{higher:+d} ulp", w, x, coded, expect) + + +# ------------------------------------------------------------------ the coded head against the BF16 fragments + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_coded_head_is_the_bf16_head_bit_for_bit(nreq) -> None: + """Integer data (fp32 sums exact in any order) with exceptions on both sides of the exponent window: the coded + head's winning keys (value and global id) and its bookkeeping are the BF16 fragments' bit for bit, on one rank and + across the ranks.""" + g = gen(seed_of(f"coded vs bf16 {nreq}")) + sign = torch.randint(0, 2, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * 2 - 1 + w = (torch.randint(1, 41, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * sign).to( + torch.bfloat16 + ) + # Exceptions on both sides of the window of 1..40: exact zeros, and 2^15 (fp32 sums stay below 2^24). + w[::5, ::512] = 0.0 + w[::13, 7::600] = 32768.0 + x = torch.randint(-2, 3, (nreq, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + wf_c, cd = _head(w, True) + assert cd[1].numel() > 0, "the shard must carry exceptions" + wf_b, cb = _head(w, False) + outs = [] + for wf, c in ((wf_c, cd), (wf_b, cb)): + args = _args(x, None, nreq, wf, 1, 1, 1, gen(seed_of("bookkeeping")), coded=c) + args = list(args) + args[9] = 0 + with validating(op): + op(*args) + torch.cuda.synchronize() + outs.append(args) + coded_args, bf16_args = outs + for i in (11, 13, 15, 16, 22): # draft_out, ids_next, pos_next, kv_lens, out_keys + assert torch.equal(coded_args[i], bf16_args[i]), i + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), DRAFT_SHARD) + assert torch.equal(coded_args[22], make_keys(best, idx)), ( + "the winning keys are not the in-tree ones" + ) + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_tiny_exception_products_are_exact(nreq) -> None: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1, and a one-hot + activation on column 0: the exceptions' exact products are the only non-zero logits, so the coded head must pick + row 1 with the BF16 head's key bit for bit (a correction against a non-zero placeholder cancels them to 0).""" + n, k = 256, 1024 + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + x = torch.zeros((nreq, k), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + wf_c, cd = _head(w, True) + assert cd[3] == 113 + keys = [] + for wf, c in ((wf_c, cd), _head(w, False)): + args = _args(x, None, nreq, wf, 0, 1, 1, gen(seed_of("tiny exceptions")), coded=c) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert args[13].tolist() == [1] * nreq + keys.append(args[22]) + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 0) + assert torch.equal(keys[0], keys[1]) and torch.equal(keys[0], make_keys(best, idx)) + + +def test_prepare_head_codes_the_shard_and_fails_closed(monkeypatch) -> None: + """With the coding on, prepare_head keeps the coded shard when it decodes to the weight bit for bit; a coding + that does not is dropped for the BF16 fragments. (The tp2 target ships HEAD_CODED = False.)""" + monkeypatch.setattr(draft_step_kernels, "HEAD_CODED", True) + g = gen(seed_of("prepare coded")) + w = randn((4096, 1024), g, 0.02) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 2) + assert prepared.coded and prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(head_weight(prepared.wf, 1024, prepared.exc_off, prepared.exc_pos, prepared.exc_w, + prepared.exp_base), w) # fmt: skip + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[11] ^= 0x10 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(draft_step_kernels, "encode_chunked", corrupt) + dropped = draft_step_kernels.prepare_head(head, 2) + assert not dropped.coded and dropped.exc_off is None + assert torch.equal(dropped.wf, bf16_fragments(w)) + + +def test_prepare_head_keeps_bf16_when_coding_is_off(monkeypatch) -> None: + """HEAD_CODED = False keeps the BF16 fragments even for a shard whose coding would decode bit for bit.""" + g = gen(seed_of("prepare bf16")) + w = randn((4096, 1024), g, 0.02) + monkeypatch.setattr(draft_step_kernels, "HEAD_CODED", False) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2) + assert not prepared.coded and prepared.exc_off is None + assert torch.equal(prepared.wf, bf16_fragments(w)) + + +def test_guard_refuses_head_rows_past_x() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[1] = args[1] + 1000 + with pytest.raises(AssertionError, match="x_rows must index x"): + with validating(op): + op(*args) + + +def test_guard_refuses_too_few_ctas() -> None: + args = list(_case(op.CELLS[1].spec, 2)) + args[20] = 8 + with pytest.raises(AssertionError, match="tiles per CTA"): + with validating(op): + op(*args) + + +# ------------------------------------------------- one pass as two launches (uGPU, kernels/m3_draft_step.UGPU) + + +@pytest.mark.parametrize( + "cell", + [c for c in op.CELLS if not c.spec["coded"]], + ids=[c.why[:44] for c in op.CELLS if not c.spec["coded"]], +) +def test_two_launches_per_domain_are_the_one_launch_pass(cell) -> None: + """A pass split into two launches (tiles [0, split) and [split, n), one per locality domain stream, the exchange + and bookkeeping by the last of both launches' CTAs) writes exactly what the one-launch pass writes: tokens, + winning keys, draft columns, next positions, KV and sequence lengths, and re-arms the counters and keys.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_locality import ( # noqa: E501 + LocalityMap, + ) + + one = _case(cell.spec, seed_of(cell.why)) + two = _case(cell.spec, seed_of(cell.why)) + op(*one) + lm = LocalityMap.get(torch.cuda.current_device()) + streams = lm.domain_streams() + n_tiles = two[3].shape[0] + split = max(1, n_tiles // 2) + per = min(lm.sms_per_domain, split, n_tiles - split) + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + for st, lo, hi in ((streams[0], 0, split), (streams[1], split, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + op(*two[:20], per, False, *two[22:], tile_lo=lo, tile_hi=hi, total_ctas=2 * per) + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) + torch.cuda.synchronize() + for i, name in ( + (11, "draft_out"), + (13, "ids_next"), + (15, "pos_next"), + (16, "kv_lens"), + (22, "out_keys"), + ): + assert torch.equal(one[i], two[i]), name + if one[18] is not None: + assert torch.equal(one[18], two[18]), "seq_lens" + assert two[6].tolist() == [0, 1] and not bool(two[5].any()), "counters / keys not re-armed" + + +# ----------------------------------------------- equal-chunk split (kernels/m3_draft_step.EQ_CHUNKS: bnd / bnd_cnt set) + + +def _eq_buffers(n_tiles: int): + return ( + torch.zeros(n_tiles, 32, 4, dtype=torch.float32, device="cuda"), + torch.zeros(n_tiles, dtype=torch.int32, device="cuda"), + ) + + +@pytest.mark.parametrize( + "cell", + [c for c in op.CELLS if not c.spec["coded"]], + ids=[c.why[:44] for c in op.CELLS if not c.spec["coded"]], +) +@pytest.mark.parametrize("launches", [1, 2]) +def test_equal_chunk_split_is_the_whole_tile_pass(cell, launches) -> None: + """Each CTA streaming an equal share of the (tile, chunk) space, the CTA-boundary tiles summed from two partials in + bnd: the same tokens, winning global ids, draft columns, positions and lengths as whole tiles per CTA (the logits + differ at most by the order of fp32 additions: their BF16 keys by at most one ulp), and bnd / bnd_cnt are left + zero for the next call, over one launch and over the two per-domain launches.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_locality import ( # noqa: E501 + LocalityMap, + ) + + whole = _case(cell.spec, seed_of(cell.why)) + eq = _case(cell.spec, seed_of(cell.why)) + n_tiles = eq[3].shape[0] + bnd, bnd_cnt = _eq_buffers(n_tiles) + op(*whole) + if launches == 1: + op(*eq, tile_lo=0, tile_hi=0, total_ctas=0, bnd=bnd, bnd_cnt=bnd_cnt) + else: + lm = LocalityMap.get(torch.cuda.current_device()) + streams = lm.domain_streams() + split = max(1, n_tiles // 2) + per = min(lm.sms_per_domain, split, n_tiles - split) + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + for st, lo, hi in ((streams[0], 0, split), (streams[1], split, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + op(*eq[:20], per, False, *eq[22:], tile_lo=lo, tile_hi=hi, total_ctas=2 * per, bnd=bnd, + bnd_cnt=bnd_cnt) # fmt: skip + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) + torch.cuda.synchronize() + for i, name in ((11, "draft_out"), (13, "ids_next"), (15, "pos_next"), (16, "kv_lens")): + assert torch.equal(whole[i], eq[i]), name + assert torch.equal(key_index(whole[22]), key_index(eq[22])), "winning global ids" + if whole[18] is not None: + assert torch.equal(whole[18], eq[18]), "seq_lens" + assert not bool(bnd.any()) and not bool(bnd_cnt.any()), "bnd / bnd_cnt not restored to zero" + assert eq[6].tolist()[0] == 0 and not bool(eq[5].any()), "counters / keys not re-armed" + + +def test_guard_refuses_an_equal_chunk_split_below_a_tile_per_cta() -> None: + """With fewer tiles than CTAs a tile could gather three partials; the op refuses the split.""" + spec = dict(op.CELLS[0].spec) + args = _case(spec, 7) + n_tiles = args[3].shape[0] + bnd, bnd_cnt = _eq_buffers(n_tiles) + with pytest.raises(RuntimeError, match="a tile per CTA"): + op( + *args[:20], + n_tiles + 1, + *args[21:], + tile_lo=0, + tile_hi=0, + total_ctas=0, + bnd=bnd, + bnd_cnt=bnd_cnt, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py new file mode 100644 index 000000000000..3c71a4a39f07 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the forward-tail spec fold (kernels/m3_fold_tail and M3SpecStepFold.in_forward). + +The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers +(step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing +when the step was not armed. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( + m3_fold_tail, + m3_spec_tail, +) + +assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" + +VOCAB = 200064 +NUM_SLOTS = 9 +MAX_ROWS = 8 +STALE = -4 + + +def _store(widths, g): + wn, wx, wd = widths + i32 = dict(dtype=torch.int32, device="cuda") + return types.SimpleNamespace( + new_tokens=torch.randint( + 0, VOCAB, (wn, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_new_tokens=torch.randint( + 0, VOCAB, (wx, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_draft_tokens=torch.randint( + 0, VOCAB, (NUM_SLOTS, wd), generator=g, device="cuda", dtype=torch.int32 + ), + new_tokens_lens=torch.randint(1, 5, (NUM_SLOTS,), generator=g, **i32), + ) + + +def _clone(ns): + return types.SimpleNamespace(**{k: v.clone() for k, v in vars(ns).items()}) + + +def _outputs(rows, t, g): + d = t - 1 + return { + "new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_draft_tokens": torch.randint( + 0, VOCAB, (rows, d), generator=g, device="cuda", dtype=torch.int32 + ), + "new_tokens_lens": torch.randint( + 1, t + 1, (rows,), generator=g, device="cuda", dtype=torch.int32 + ), + } + + +def _engine(t): + i32 = dict(dtype=torch.int32, device="cuda") + return ( + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS * (t - 1),), STALE, **i32), + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS,), STALE, **i32), + ) + + +def _reference(store, engine, slots, out, t): + """slot_scatter into the store, then step_gather of the next step's inputs from it.""" + assert m3_spec_tail.slot_scatter( + slots, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["new_tokens_lens"], + store.new_tokens_lens, + out["next_draft_tokens"], + store.next_draft_tokens, + ) + p = slots.numel() + ids, draft, pos_off, kv_off = engine + slots32 = slots.to(torch.int32) + assert m3_spec_tail.step_gather( + store.next_new_tokens, + store.next_draft_tokens, + store.new_tokens_lens, + slots32, + slots32.repeat_interleave(t), + p, + t, + t - 1, + ids, + 0, + draft, + 0, + pos_off, + kv_off, + 0, + ) + + +def _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=False): + ids, draft, pos_off, kv_off = engine + assert m3_fold_tail.fold_tail( + ctrl, + slots_buf, + rows, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["next_draft_tokens"], + store.next_draft_tokens, + out["new_tokens_lens"], + store.new_tokens_lens, + ids, + draft, + pos_off, + kv_off, + t, + cpp=cpp, + ) + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("widths", [(4, 4, 3), (8, 8, 7)], ids=["store=step", "store wider"]) +def test_tail_matches_scatter_then_gather(rows, widths, cpp) -> None: + """Store bit for bit; engine rows [0, P*T) / [0, P) bit for bit (the forward reads only those rows).""" + t = 4 + g = torch.Generator(device="cuda").manual_seed(1000 * rows + widths[0]) + store = _store(widths, g) + out = _outputs(rows, t, g) + slots = torch.randperm(NUM_SLOTS, generator=g, device="cuda")[:rows] + ref_store, tail_store = _clone(store), _clone(store) + ref_engine, tail_engine = _engine(t), _engine(t) + _reference(ref_store, ref_engine, slots, out, t) + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + slots_buf = torch.zeros(MAX_ROWS, dtype=torch.long, device="cuda") + slots_buf[:rows] = slots + _tail(tail_store, tail_engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(ref_store, name), getattr(tail_store, name)), name + spans = (rows * t, rows * (t - 1), rows * t, rows) + for k, (ref, got, n) in enumerate(zip(ref_engine, tail_engine, spans)): + assert torch.equal(ref[:n], got[:n]), k + assert torch.all(got[n:] == STALE), k + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +def test_unarmed_tail_writes_nothing(cpp) -> None: + """ctrl[0] = 0 (the step was not armed): the store and the engine buffers are untouched.""" + t, rows = 4, 2 + g = torch.Generator(device="cuda").manual_seed(7) + store = _store((4, 4, 3), g) + out = _outputs(rows, t, g) + before, engine = _clone(store), _engine(t) + ctrl = torch.zeros(1, dtype=torch.int32, device="cuda") + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(before, name), getattr(store, name)), name + for buf in engine: + assert torch.all(buf == STALE) + + +def _fold(): + fold = modeling.M3SpecStepFold() + t = 4 + engine = _engine(t) + fold.bind_engine_buffers(*engine, max_batch_size=MAX_ROWS) + g = torch.Generator(device="cuda").manual_seed(11) + fold.bind_sampler_store(_store((4, 4, 3), g)) + return fold, g, t + + +def test_fold_handshake_stages_only_confirmed_rows() -> None: + """begin_step arms; the sampler confirms only the armed rows from the tail's outputs; the next step is staged only + for exactly those rows in that order.""" + fold, g, t = _fold() + assert fold.begin_step([3, 5], 2) is False # nothing staged before the first forward + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [3, 5]) is True + assert fold.begin_step([3, 5], 2) is True + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert ( + fold.sampler_rows_stored(out, 0, [5, 3]) is False + ) # other order: the sampler stores itself + assert fold.begin_step([5, 3], 2) is False + + +def test_fold_handshake_needs_the_tail_outputs() -> None: + """Outputs the tail did not run on (another dict, a mixed step, an unarmed step) are never reported stored.""" + fold, g, t = _fold() + fold.begin_step([1], 1) + fold.store_in_forward(_outputs(1, t, g)) + assert fold.sampler_rows_stored(_outputs(1, t, g), 0, [1]) is False + fold.begin_step(None, 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [1]) is False + fold.begin_step([1], 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 1, [1]) is False # rows were skipped (chunked context) + + +def test_a_padded_step_after_an_armed_one_writes_nothing() -> None: + """K3 #150 check #5: CUDA-graph padding rows carry no seq slot, so the engine passes rows=None for a padded step. + Even right after an armed step, that step uploads ctrl 0 and its forward tail leaves the store and the engine + buffers untouched (a dummy row never writes a live slot).""" + fold, g, t = _fold() + fold.begin_step([2, 4], 2) + fold.store_in_forward(_outputs(2, t, g)) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 2 + + fold.begin_step(None, 4) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 0 + store_before = _clone(fold.store) + engine_before = [buf.clone() for buf in fold.step_inputs] + fold.store_in_forward(_outputs(4, t, g)) + torch.cuda.synchronize() + for name in vars(store_before): + assert torch.equal(getattr(store_before, name), getattr(fold.store, name)), name + for before, buf in zip(engine_before, fold.step_inputs): + assert torch.equal(before, buf) + + +def test_the_fold_stores_in_the_forward_once_bound() -> None: + """The fold always stores and stages at the forward's tail; before the engine binds its buffers nothing is armed.""" + fold = modeling.M3SpecStepFold() + assert fold.in_forward + assert fold.begin_step([0], 1) is False + + +def test_the_store_copy_event_is_recorded_at_binding(): + """An external event the sampler records after its side-stream store copies and the forward's tail waits on, + recorded once at binding so a captured wait has a completed record to wait on.""" + fold = modeling.M3SpecStepFold() + device = torch.device("cuda") + fold.bind_engine_buffers( + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + 4, + ) + assert fold.store_copies_done is not None + fold.store_copies_done.synchronize() + + +def test_the_store_wait_precedes_only_a_plain_launch_under_capture(): + """The store wait is an external-event wait node under graph capture, which drops the next kernel's PDL edge: + it may precede the fold tail (a plain launch); any other launch fails the capture (eager steps are not checked).""" + fold = modeling.M3SpecStepFold() + device = torch.device("cuda") + fold.bind_engine_buffers( + *(torch.zeros(8, dtype=torch.int32, device=device) for _ in range(4)), 4 + ) + assert fold.STORE_WAIT_BEFORE == frozenset({"fold_tail"}) + fold._wait_store_copies(before="draft_step") # eager: allowed + x = torch.zeros(4, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + fold._wait_store_copies(before="fold_tail") + x.add_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(x, torch.ones(4, device=device)) + bad = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="PDL edge"): + with torch.cuda.graph(bad, stream=stream): + fold._wait_store_copies(before="draft_step") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py new file mode 100644 index 000000000000..c102a3374bc1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the per-request forced-acceptance draw (``FORCE_AL_PER_REQUEST``, the ``m3_force_vec_req`` entry, +synth only): each request's forced count comes from its own key (its KV length at the verify), so it does not depend +on the call history, on the request's slot in the batch, or on the other requests.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec_req, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SEED, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _pool(): + g = torch.Generator(device="cpu") + g.manual_seed(_FORCE_ACCEPT_RNG_SEED) + return torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, dtype=torch.float32, generator=g).cuda() + + +def _draw(pool, keys, base_total=2, max_total=4, frac=0.6): + out = torch.empty(keys.numel(), dtype=torch.int32, device="cuda") + m3_force_vec_req( + pool, + keys, + out, + keys.numel(), + base_total, + max_total, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (keys.numel() - 1).bit_length()), + ) + return out + + +def test_a_request_draws_the_same_whatever_ran_before_or_beside_it() -> None: + pool = _pool() + keys = torch.tensor([165_521, 40_003, 323_007, 7], dtype=torch.int32, device="cuda") + first = _draw(pool, keys) + for _ in range(5): # no shared counter: repeated calls draw the same + assert torch.equal(_draw(pool, keys), first) + perm = torch.tensor([2, 0, 3, 1], device="cuda") + assert torch.equal(_draw(pool, keys[perm]), first[perm]) # slot-independent + alone = torch.cat([_draw(pool, keys[i : i + 1]) for i in range(4)]) + assert torch.equal(alone, first) # independent of the other requests + + +def test_the_counts_follow_the_forced_fraction() -> None: + pool = _pool() + keys = torch.arange(100_000, 100_000 + 64 * 400, dtype=torch.int32, device="cuda") + counts = torch.cat([_draw(pool, keys[i : i + 64]) for i in range(0, keys.numel(), 64)]) + assert set(counts.tolist()) <= {2, 3} + extra = (counts == 3).float().mean().item() + assert abs(extra - 0.6) < 0.02, extra diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..3bad6ccc5d2a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec entry of the MiniMax-M3 tp2 target: the fractional forced-acceptance draw.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(5) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + counter = torch.tensor([41], dtype=torch.int64, device="cuda") + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + counter, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert int(counter.item()) == 42, "the draw bumps the step counter once" + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py new file mode 100644 index 000000000000..53a58df60bef --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec_req entry of the MiniMax-M3 tp2 target: the per-request fractional forced-acceptance +draw (one draw per request at its own key, no shared counter).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec_req as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(7) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + # Each request's key is its KV length at the verify: long, distinct, arbitrary. + keys = torch.randint( + 1, 1 << 20, (spec["batch"],), generator=g, device="cuda", dtype=torch.int64 + ).to(torch.int32) + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + keys, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py new file mode 100644 index 000000000000..e68903313514 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (kernels/m3_lowm_splitk, +modeling.DRAFT_LOWM_OWN_SPLITK): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's +linears, eager and in a CUDA graph, with and without PDL.""" + +import pytest +import torch +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + tactic_for, +) + +assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" + +# Every certified (N, K, rows) of the TP2 drafter table that a split-K tactic serves: (N, K, rows, tactic). +CELLS = [ + (nk[0], nk[1], m, tactic_for(nk[0], nk[1], m)[1]) + for nk, rows in sorted(ROWS.items()) + for m in rows + if tactic_for(nk[0], nk[1], m)[0] == "s" +] + + +def test_the_table_has_the_certified_split_k_cells() -> None: + assert sorted((n, k, m) for n, k, m, _ in CELLS) == [ + (6144, 4096, 4), + (6144, 4096, 8), + (6144, 4096, 16), + (6144, 9216, 4), + (6144, 9216, 16), + (6144, 9216, 32), + (6144, 18432, 16), + ] + + +@pytest.mark.parametrize("use_pdl", [False, True], ids=["plain", "pdl"]) +@pytest.mark.parametrize("cell", CELLS, ids=[f"{n}x{k}@{m}" for n, k, m, _ in CELLS]) +def test_bitwise_flashinfer(cell, use_pdl) -> None: + n, k, m, tactic = cell + g = torch.Generator(device="cuda").manual_seed(n + 7 * k + m) + for trial in range(3): + x = (torch.randn(m, k, generator=g, device="cuda") * (1.0 + trial)).to(torch.bfloat16) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + ref = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + got = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + fi_run(x, w.t(), None, ref, use_pdl, FiTactic(*tactic)) + run_splitk_dense(x, w.t(), None, got, use_pdl, SplitKTactic(*tactic)) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), f"trial {trial}" + # In a CUDA graph, as the drafter runs it. + cap = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_splitk_dense(x, w.t(), None, cap, use_pdl, SplitKTactic(*tactic)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(cap.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..2af99c2062f6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp2 m3_sampled_accept entry (fused temperature + top-p verify over the 100032-column TP2 shard) +and its workspace size.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_sampled_accept import ( # noqa: E501 + m3_sampled_accept as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_sampled_accept_ws_words import ( # noqa: E501 + m3_sampled_accept_ws_words as ws_words, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" + +DRAFT_LEN = 3 +#: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). +HEAD = 12 +#: Draws per distribution cell, and the total-variation band against the reference distribution. With 12 +#: categories and >= 4000 draws the expected TV of an exact sampler is ~0.02 (sd ~0.006); 0.06 is ~6 sd above it +#: while a wrong nucleus cut (one category too many or too few, >= 1 / 12 of the mass moved) lands at >= 0.08. +DRAWS = 4000 +TV_BAND = 0.06 +#: The cells of one kernel launch (<= 48 verify rows; every certified cell): the tests below call the kernel itself. +SINGLE = [c for c in op.CELLS if c.spec["nreq"] * (DRAFT_LEN + 1) <= 48] + + +def _logits(spec, g, rows): + vs = spec["vs"] + if spec["check"] == "flat": + return (torch.randn((rows, vs), generator=g, device="cuda") * 0.3).to(torch.bfloat16) + if spec["check"] == "argmax": + return (torch.randn((rows, vs), generator=g, device="cuda") * 3).to(torch.bfloat16) + row = torch.full((vs,), -20.0, device="cuda") + hot = torch.randperm(vs, generator=g, device="cuda")[:HEAD] + row[hot] = torch.linspace(2.0, 0.2, HEAD, device="cuda") + return row.to(torch.bfloat16).expand(rows, vs).contiguous() + + +def _case(spec, seed): + g = gen(seed) + nreq = spec["nreq"] + rows = nreq * (DRAFT_LEN + 1) + logits = _logits(spec, g, rows) + words = ws_words(1) + uc = torch.zeros(words, dtype=torch.int64, device="cuda") + counters = torch.zeros(2, dtype=torch.int32, device="cuda") + tokens = torch.zeros(max(48, rows), dtype=torch.int32, device="cuda") + extra = torch.zeros(1, dtype=torch.int32, device="cuda") + temperatures = torch.full((rows,), spec["temperature"], dtype=torch.float32, device="cuda") + top_ps = torch.full((rows,), spec["top_p"], dtype=torch.float32, device="cuda") + seeds = torch.randint(0, 2**62, (rows,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.zeros(rows, dtype=torch.int64, device="cuda") + # Drafts drawn from the likeliest tokens, so some prefixes match. + top = logits[0].float().topk(3).indices.to(torch.int32) + spec["vocab_offset"] + draft = top[torch.randint(0, 3, (nreq, DRAFT_LEN), generator=g, device="cuda")].contiguous() + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + force = spec.get("force", "none") + force_total = 3 if force == "total" else 0 + force_vec = (torch.randint(1, DRAFT_LEN + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if force == "vec" else None) # fmt: skip + return [logits, spec["vocab_offset"], temperatures, top_ps, seeds, offsets, uc, uc.data_ptr(), counters, tokens, + 0, 1, draft, batch_idx, acc, num, gather, force_total, next_new, force_vec, extra, True] # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Statistical for the draw, exact for everything derived from it. + + Every drawn token has nonzero reference probability (lies in the nucleus); the acceptance outputs equal the + reference's strict acceptance of the drawn tokens, every call; the call counter advances once per call; and, + for the distribution cells, the draws over many calls (fresh Philox offsets) are within TV_BAND of the + reference distribution (see DRAWS / TV_BAND). + """ + spec = cell.spec + args = _case(spec, seed_of(cell.why)) + logits, offset, rows = args[0], spec["vocab_offset"], args[0].shape[0] + certifying(op) + ref = op.reference(*args) + probs = ref["probs"] + calls = 1 if spec["check"] == "argmax" else -(-DRAWS // rows) + counts = torch.zeros(logits.shape[1], dtype=torch.float64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + drawn = args[9][:rows].clone() + local = drawn.long() - offset + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + assert bool((probs.gather(1, local[:, None]) > 0).all()), ( + "a drawn token outside the nucleus" + ) + acc = ref["accept"](drawn) + assert torch.equal(args[14], acc["accepted"]) + assert torch.equal(args[15], acc["num_accepted"]) + assert torch.equal(args[16], acc["gather_ids"]) + assert torch.equal(args[18], acc["next_new"]) + # One launch per call (every cell is <= 48 rows). + assert args[8].tolist() == [0, (c + 1) * -(-rows // 48)] + counts.index_add_(0, local, torch.ones(rows, dtype=torch.float64, device="cuda")) + if spec["check"] == "distribution": + emp = counts / counts.sum() + tv = 0.5 * (emp - probs[0].double()).abs().sum().item() + assert tv <= TV_BAND, f"total variation {tv:.3f} against the reference nucleus distribution" + + +@pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) +def test_prewait_parameter_loads_change_nothing(cell) -> None: + """kernels/m3_sampled_accept.PREWAIT_PARAMS moves only where the sampling parameters load (before or after the PDL + wait): the same draws and acceptance outputs, bit for bit, call after call.""" + outs = [] + for prewait in (True, False): + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + got = [] + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept(*args, prewait) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + outs.append(got) + for call, (a, b) in enumerate(zip(*outs)): + for x, y in zip(a, b): + assert torch.equal(x, y), f"call {call}: prewait changes {x.tolist()} -> {y.tolist()}" + + +@pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) +def test_the_early_trigger_changes_nothing(cell) -> None: + """kernels/m3_sampled_accept.EARLY_TRIGGER moves only where the grid releases its dependents: the same draws and + acceptance outputs, bit for bit, call after call, PDL launch on.""" + outs = [] + for early in (True, False): + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + got = [] + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept(*args, True, early) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + outs.append(got) + for call, (a, b) in enumerate(zip(*outs)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the early trigger changes {x.tolist()} -> {y.tolist()}" + ) + + +def _walk_outputs(args, ilp, calls=3, offset_step=1000): + """The draws and acceptance outputs of `calls` calls with the owner walk one token (ilp False) or 8 tokens at a + time (ilp True), each call at its own Philox offsets.""" + rows = args[0].shape[0] + got = [] + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * offset_step) + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept(*args, True, True, ilp) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + return got + + +@pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) +def test_the_ilp_walk_changes_nothing(cell) -> None: + """kernels/m3_sampled_accept.ILP_WALK (SA-PW) only computes the owner's token weights 8 at a time; the walk adds + them in the same order with the same early exit, so the draws and acceptance outputs are the one-token walk's, + bit for bit, call after call.""" + one = _walk_outputs(_case(cell.spec, seed_of(cell.why)), False) + ilp = _walk_outputs(_case(cell.spec, seed_of(cell.why)), True) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.95]) +def test_the_ilp_walk_changes_nothing_over_many_rounds(top_p) -> None: + """A flat tail (many tokens of equal weight) under a tiny and a wide nucleus: rows take many exchange rounds and + each round's walk exits at a different token; the 8-token walk draws the same tokens as the one-token walk.""" + spec = dict(nreq=1, vs=100032, vocab_offset=100032, check="flat", top_p=top_p, temperature=1.0) + one = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), False, calls=24) + ilp = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), True, calls=24) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +def test_the_ilp_walk_is_the_same_under_graph_replay() -> None: + """The 8-token walk captured in a CUDA graph and replayed draws what the one-token walk draws eagerly.""" + cell = op.CELLS[0] + eager = _walk_outputs(_case(cell.spec, seed_of(cell.why)), False) + _walk_outputs( + _case(cell.spec, seed_of("ilp walk warm-up")), True, calls=1 + ) # first launch outside the capture + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept(*args, True, True, True) + torch.cuda.current_stream().wait_stream(stream) + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + graph.replay() + torch.cuda.synchronize() + for x, y in zip(eager[c], [args[i] for i in (9, 14, 15, 16, 18)]): + assert torch.equal(x, y), ( + f"replay {c}: the graphed 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.0002]) +def test_long_rejection_chains_stay_in_the_nucleus(top_p) -> None: + """A tiny nucleus over a flat tail: rows that need more exchange rounds than the workspace has slots. + + Each round's candidate is drawn over the tokens heavier than the last one, so a row takes ~ln(1 / top_p) + 1 + rounds on average, with a tail past the 16 exchange slots (about 1% of the rows at top_p 0.001). Rounds reuse the + slots cyclically, told apart by a lap-parity bit in the wire tags, and no round accepts unconditionally: every + draw, however many rounds it took, lies in the reference nucleus. World 1 (one rank's 8 CTAs per row exchange + through the same slot ring); the cross-rank path needs the multi-device job. + """ + spec = dict(nreq=1, vs=100032, vocab_offset=100032, check="flat", top_p=top_p, temperature=1.0) + args = _case(spec, seed_of(f"flat tail top_p {top_p}")) + logits, rows = args[0], args[0].shape[0] + probs = op.reference(*args)["probs"] + calls = -(-DRAWS // rows) + outside = torch.zeros((), dtype=torch.int64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + local = args[9][:rows].long() - spec["vocab_offset"] + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + outside += (probs.gather(1, local[:, None]) == 0).sum() + assert args[8].tolist() == [0, c + 1] + draws = calls * rows + rounds = 2.0 + args[20].item() / draws + kept = int((probs > 0).sum(1).min().item()) + assert outside.item() == 0, ( + f"{outside.item()} of {draws} draws outside the nucleus (kept >= {kept} of {logits.shape[1]}, " + f"{rounds:.2f} rounds per row on average)" + ) + assert rounds > 4.0, ( + f"only {rounds:.2f} rounds per row: the rows no longer stress the slot ring" + ) + + +@pytest.mark.parametrize("cell", ws_words.CELLS, ids=[c.why[:44] for c in ws_words.CELLS]) +def test_workspace_words(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(ws_words) + world = cell.spec["world"] + assert ws_words(world) == ws_words.reference(world) + + +def test_guard_refuses_top_k_rows() -> None: + args = _case(op.CELLS[0].spec, 1) + top_ks = torch.full((args[0].shape[0],), 20, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="top-k rows"): + with validating(op): + op(*args, top_ks=top_ks) + + +def test_guard_refuses_fp64_temperatures() -> None: + args = _case(op.CELLS[0].spec, 2) + args[2] = args[2].double() + with pytest.raises(AssertionError, match="fp32 sampling parameters"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..5c6a99bb495a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_shard_accept entry (greedy verify from this rank's logits shard).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + craft_drafts, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" + + +def _case(spec, seed, edit_logits=None): + g = gen(seed) + nreq, d, vs, world, rank = ( + spec["nreq"], + spec["draft_len"], + spec["vs"], + spec["world"], + spec["rank"], + ) + m = nreq * (d + 1) + if spec["strided"]: + logits = randn((m, vs + 72), g, 3.0)[:, :vs] + else: + logits = randn((m, vs), g, 3.0) + if edit_logits is not None: + edit_logits(logits) + uc, keys, counters = head_workspace(world, slot_rows=SHARD_ROWS) + own = logits.float().max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=SHARD_ROWS, ties=spec.get("ties", False) + ) + offset = rank * vs + best, idx = local_argmax(logits.float(), offset) + if world == 1: + tokens = idx.to( + torch.int32 + ) # no peers: this rank's own pick (any row count, incl. W2's chunked calls) + else: + tokens = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, SHARD_ROWS) + draft = craft_drafts(tokens, nreq, d) + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + force_total = 3 if spec["force"] == "total" else 0 + force_vec = (torch.randint(1, d + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if spec["force"] == "vec" else None) # fmt: skip + return (logits, offset, keys, counters, uc, uc.data_ptr(), rank, world, draft, batch_idx, acc, num, gather, + force_total, next_new, num_sms(), force_vec) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bit-exact: the op reduces BF16 logits without arithmetic, ties (within the shard and across ranks) go to + the lowest global index on both sides.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) + # One call per launch: a chunked (W2, > 48-row) call launches once per 12-request chunk. + launches = -(-args[0].shape[0] // SHARD_ROWS) + assert args[3].tolist() == [0, launches] and not bool(args[2].any()) + + +def test_guard_refuses_nan_logits() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0][1, 7] = float("nan") + with pytest.raises(AssertionError, match="NaN"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_batch_index_past_the_requests() -> None: + args = list(_case(op.CELLS[0].spec, 2)) + args[9] = args[9] + 1 + with pytest.raises(AssertionError, match="batch_idx"): + with validating(op): + op(*args) + + +def _signed_zero_maxima(logits: torch.Tensor) -> None: + """Every row's maximum is zero, reached as -0.0 at a lower index and +0.0 at a higher one.""" + logits.fill_(-1.0) + logits[:, 3] = -0.0 + logits[:, 9] = 0.0 + + +@pytest.mark.parametrize("ties", [False, True], ids=["peer_rows", "peer_ties"]) +def test_signed_zero_maxima_take_the_lower_index(ties) -> None: + """-0.0 and +0.0 compare equal, so the lower index wins as with torch.argmax (also across ranks).""" + spec = dict(op.CELLS[1].spec, ties=ties) + args = _case(spec, seed_of(f"signed-zero-{ties}"), edit_logits=_signed_zero_maxima) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..f0417a37a805 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_slot_scatter entry (the spec sampler's per-slot stores in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_slot_scatter import ( # noqa: E501 + m3_slot_scatter as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" + +STALE = -4 + + +def _case(spec, seed): + """Worker outputs and the store's transposed [slots, width] views of its [width, slots, 1] buffers.""" + g = gen(seed) + n, nslots = spec["n"], spec["slots"] + + def rnd(shape): + return torch.randint(0, VOCAB, shape, generator=g, device="cuda", dtype=torch.int32) + + store_new = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_next = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_lens = torch.full((nslots,), STALE, dtype=torch.int32, device="cuda") + store_draft = torch.full((nslots, spec["store_d"]), STALE, dtype=torch.int32, device="cuda") + slots = torch.randperm(nslots, generator=g, device="cuda")[:n] + lens = torch.randint(1, 5, (n,), generator=g, device="cuda", dtype=torch.int32) + return (slots, rnd((n, spec["src_w"])), store_new.squeeze(-1).T, rnd((n, spec["src_w"])), + store_next.squeeze(-1).T, lens, store_lens, rnd((n, spec["src_d"])), store_draft) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, zero padding included, other slots untouched.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in (("da", args[2]), ("db", args[4]), ("dl", args[6]), ("dd", args[8])): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_repeated_slots() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0] = args[0].clone() + args[0][1] = args[0][0] + with pytest.raises(AssertionError, match="distinct"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py new file mode 100644 index 000000000000..b023cbd47add --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for m3_spec_offsets (RuntimeHooks.spec_offsets_op): the overlap scheduler's two speculative step corrections +in one plain launch, bit for bit the in-place ATen adds, eager and in a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_spec_offsets import ( # noqa: E501 + m3_spec_offsets as op, +) + +assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" + +I32 = torch.iinfo(torch.int32) + + +def _inputs(tokens, seqs, seed, wrap=False): + g = torch.Generator(device="cuda").manual_seed(seed) + lo, hi = (I32.min, I32.max) if wrap else (0, 1 << 20) + pos = torch.randint(lo, hi, (tokens,), generator=g, device="cuda", dtype=torch.int32) + kv = torch.randint(lo, hi, (seqs,), generator=g, device="cuda", dtype=torch.int32) + pos_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (tokens,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + kv_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (seqs,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + if wrap: + pos[:2], pos_off[:2] = I32.max, 1 # overflow wraps, as the ATen add does + kv[:1], kv_off[:1] = I32.min, -1 + return pos, pos_off, kv, kv_off + + +CELLS = [c.spec for c in op.CELLS] + [ + dict(tokens=1, seqs=1), + dict(tokens=300, seqs=75), + dict(tokens=4, seqs=0), +] + + +@pytest.mark.parametrize("spec", CELLS, ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items())) +def test_bitwise_aten_adds(spec) -> None: + for trial in range(3): + pos, pos_off, kv, kv_off = _inputs( + spec["tokens"], spec["seqs"], 97 * trial + spec["tokens"], spec.get("wrap", False) + ) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + assert op(pos, pos_off, kv, kv_off) is True + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"trial {trial}" + + +def test_graph_replays_read_the_live_offsets() -> None: + pos, pos_off, kv, kv_off = _inputs(32, 8, 5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + for step in range(3): + fresh = _inputs(32, 8, 11 + step) + for dst, src in zip((pos, pos_off, kv, kv_off), fresh): + dst.copy_(src) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"step {step}" + + +def test_runner_views_are_accepted() -> None: + """The runner passes position_ids[0, ctx:] and kv_lens_cuda[ctx_reqs:seqs]: 1-D slices of contiguous buffers.""" + position_ids = torch.arange(40, dtype=torch.int32, device="cuda").view(1, 40) + kv_lens = torch.arange(12, dtype=torch.int32, device="cuda") + pos_off = torch.full((40,), 2, dtype=torch.int32, device="cuda") + kv_off = torch.full((12,), 3, dtype=torch.int32, device="cuda") + assert op(position_ids[0, 8:], pos_off[:32], kv_lens[2:10], kv_off[:8]) is True + torch.cuda.synchronize() + assert torch.equal( + position_ids[0], torch.cat([torch.arange(8), torch.arange(8, 40) + 2]).int().cuda() + ) + assert torch.equal( + kv_lens, + torch.cat([torch.arange(2), torch.arange(2, 10) + 3, torch.arange(10, 12)]).int().cuda(), + ) + + +@pytest.mark.parametrize( + "bad", + ["int64", "strided", "length", "2d"], +) +def test_declines_other_layouts_without_writing(bad) -> None: + pos, pos_off, kv, kv_off = _inputs(8, 2, 3) + if bad == "int64": + kv = kv.long() + elif bad == "strided": + pos = torch.zeros(16, dtype=torch.int32, device="cuda")[::2] + elif bad == "length": + pos_off = pos_off[:7] + else: + pos = pos.view(2, 4) + before = (pos.clone(), kv.clone()) + assert op(pos, pos_off, kv, kv_off) is False + torch.cuda.synchronize() + assert torch.equal(pos, before[0]) and torch.equal(kv, before[1]) + + +def test_launches_plainly() -> None: + """It writes step inputs that pre-wait readers take before their grid wait (MSA contract, dispatch + PRODUCER_PREWAIT_STEP_INPUTS): its kernel node carries no programmatic (PDL) serialization.""" + from cuda.bindings import runtime as rt + + pos, pos_off, kv, kv_off = _inputs(32, 8, 7) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + err, _, n = rt.cudaGraphGetNodes(raw, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + flags = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind == rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + flags.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + assert flags == [0], flags diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..e6e732174a8b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_gather entry (the overlap scheduler's step-input gathers in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_step_gather import ( # noqa: E501 + m3_step_gather as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" + +#: Rows of the engine's flat buffers (the max batch), and their stale fill. +MAX_ROWS = 8 +STALE = -4 + + +def _case(spec, seed): + g = gen(seed) + p, t, d, nslots, b0 = spec["p"], spec["t"], spec["d"], spec["slots"], spec["b0"] + i32 = dict(dtype=torch.int32, device="cuda") + new_tok = torch.randint(0, VOCAB, (t, nslots, 1), generator=g, device="cuda", dtype=torch.int32) + if spec.get("two_d"): + new_tok = new_tok.squeeze(-1) + next_draft = torch.randint(0, VOCAB, (nslots, d), generator=g, device="cuda", dtype=torch.int32) + lens = torch.randint(1, t + 1, (nslots,), generator=g, device="cuda", dtype=torch.int32) + slots = torch.randperm(nslots, generator=g, device="cuda")[:p].to(torch.int32) + ppi = slots.repeat_interleave(t) + ids = torch.full((MAX_ROWS * t,), STALE, **i32) + draft = torch.full((MAX_ROWS * d,), STALE, **i32) + pos_off = torch.full((MAX_ROWS * t,), STALE, **i32) + kv_off = torch.full((MAX_ROWS,), STALE, **i32) + return ( + new_tok, + next_draft, + lens, + slots, + ppi, + p, + t, + d, + ids, + spec["ids0"], + draft, + spec["d0"], + pos_off, + kv_off, + b0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in ( + ("ids", args[8]), + ("draft", args[10]), + ("pos_off", args[12]), + ("kv_off", args[13]), + ): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_a_slot_past_the_store() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[3] = args[3].clone() + args[3][0] = 8 + with pytest.raises(AssertionError, match="slots must index the store"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py new file mode 100644 index 000000000000..13fc115c0afa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_prologue entry (the decode graph's step-input prologue in one launch).""" + +import random +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_step_prologue import ( # noqa: E501 + m3_step_prologue as op, +) +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" + +MAX_DESC = StepInputPack.MAX_DESC +WINDOW = StepInputPack.GRAPH_WINDOW +DATA_START = StepInputPack.HEADER_BYTES + StepInputPack.DESC_BYTES * MAX_DESC +ARENA = 1 << 20 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _layout(spec, seed): + """Copies as (dst, src, bytes) arena offsets (a host copy's src: a window offset) and the host copies' bytes.""" + rng = random.Random(seed) + h2d, d2d, data = [], [], [] + used = DATA_START + pos = 0 + for _ in range(spec["h2d"]): + n = rng.choice([1, 3, 4, 7, 16, 24, 32, 48, 64, 5, 12]) + pos += rng.randint(0, 40) # gaps and every alignment + h2d.append((pos, used, n)) + data.append(bytes(rng.getrandbits(8) for _ in range(n))) + pos += n + used += -(-n // 16) * 16 + assert pos < ARENA // 2 and used <= WINDOW + # Device copies write the last quarter; unchained ones read the third. + dpos, spos = 3 * ARENA // 4, ARENA // 2 + for i in range(spec["d2d"]): + if spec.get("chained") and h2d: + dst_h, _, n_h = h2d[i % len(h2d)] + off = rng.randint(0, n_h - 1) + src, n = dst_h + off, n_h - off + else: + n = rng.choice([4, 8, 16, 32, 6, 64]) + spos += rng.randint(0, 24) + src = spos + spos += n + dpos += rng.randint(0, 24) + d2d.append((dpos, src, n)) + dpos += n + assert dpos < ARENA and spos < 3 * ARENA // 4 + return h2d, d2d, data + + +def _window(layout, base: int) -> torch.Tensor: + """The window of ``layout`` for an arena at device address ``base``.""" + h2d, d2d, data = layout + w = torch.zeros(WINDOW, dtype=torch.uint8) + words = [len(h2d), len(d2d)] + for dst, src, n in h2d: + words += [base + dst, src, n] + for dst, src, n in d2d: + words += [base + dst, base + src, n] + w[: 8 * len(words)] = torch.tensor(words, dtype=torch.int64).view(torch.uint8) + for (_, off, n), b in zip(h2d, data): + w[off : off + n] = torch.frombuffer(bytearray(b), dtype=torch.uint8) + return w.cuda() + + +def _arenas(seed): + a = torch.randint(0, 256, (ARENA,), generator=gen(seed), device="cuda", dtype=torch.int32) + a = a.to(torch.uint8) + return a, a.clone() + + +@pytest.mark.parametrize("with_epoch", [True, False], ids=["epoch", "no_epoch"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, with_epoch, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Byte for byte the default prologue's destinations, header clear and epoch.""" + seed = seed_of(cell.why) + layout = _layout(cell.spec, seed) + arena, ref_arena = _arenas(seed) + window, ref_window = _window(layout, arena.data_ptr()), _window(layout, ref_arena.data_ptr()) + epoch = torch.full((1,), 41, dtype=torch.int64, device="cuda") if with_epoch else None + ref_epoch = epoch.clone() if with_epoch else None + certifying(op) + op(window, epoch, MAX_DESC) + op.reference(ref_window, ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena) + # Each window holds its own arena's addresses: past the consumed header, both are left as staged. + for w, a in ((window, arena), (ref_window, ref_arena)): + assert not w[:16].any(), "the header is consumed" + assert torch.equal(w[16:], _window(layout, a.data_ptr())[16:]) + if with_epoch: + assert int(epoch.item()) == int(ref_epoch.item()) == 42 + + +def test_graph_replays_scatter_each_new_window() -> None: + """Captured once on a static window: every replay writes that replay's window and bumps the epoch once.""" + arena, ref_arena = _arenas(7) + window = torch.zeros(WINDOW, dtype=torch.uint8, device="cuda") + epoch = torch.zeros(1, dtype=torch.int64, device="cuda") + ref_epoch = epoch.clone() + op(window, epoch, MAX_DESC) # the empty window: writes nothing + op.reference(window.clone(), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(window, epoch, MAX_DESC) + for step, spec in enumerate( + [dict(h2d=24, d2d=8, chained=True), dict(h2d=0, d2d=0), dict(h2d=40, d2d=4)] + ): + layout = _layout(spec, 100 + step) + window.copy_(_window(layout, arena.data_ptr())) + graph.replay() + op.reference(_window(layout, ref_arena.data_ptr()), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena), step + assert not window[:16].any(), step + assert int(epoch.item()) == int(ref_epoch.item()) == step + 2, step + + +def _pack_case(prologue_op): + pack = StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=True, + in_graph=True, + early_copy=True, + step_epoch=True, + prologue_op=prologue_op, + ) + a, b, c = (torch.zeros(16, dtype=torch.int32, device="cuda") for _ in range(3)) + out = torch.zeros(48, dtype=torch.int32, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _pack_step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + + def host(v): + return torch.tensor(v, dtype=torch.int32).pin_memory() + + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, host(values[0])) + step_inputs.copy_(b, host(values[1])) + if overlap: + step_inputs.copy_(a[:4], host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if not flush: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_pack_with_the_one_launch_prologue_matches_the_default() -> None: + """The same pack steps through the default prologue and the one-launch one: the same outputs and epochs, and the + one-launch prologue is engaged (the default's warm-up compile is not what bumped its epoch).""" + packs = [_pack_case(None), _pack_case(op)] + assert packs[0][0].prologue_op is None and packs[1][0].prologue_op is op + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + outs = [] + for pack, graph, dsts, out in packs: + # A replay with a new pack, one with an eagerly issued pack, and one with nothing to scatter. + _pack_step(pack, graph, dsts, (va, vb), overlap=step == 2, flush=step != 3) + torch.cuda.synchronize() + outs.append((out.tolist(), int(pack.epoch.item()))) + assert outs[0] == outs[1], step + assert outs[1][1] == step + 2, step # one bump at warm-up, then one per replay diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py new file mode 100644 index 000000000000..f89a2ddde2e6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""W2b (SpecTail, wide target): the in-tree draft loop drafts from the drafter head's local vocab shard.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.modeling import ( # noqa: E501 + M3Eagle3Worker, + _M3ShardDraftLogits, +) +from tensorrt_llm._torch.modules.linear import TensorParallelMode + +pytestmark = pytest.mark.cpu_only + + +class _Head(nn.Module): + tp_mode = TensorParallelMode.COLUMN + padding_size = 0 + tp_size = 2 + bias = None + + def apply_linear(self, x, bias): + assert bias is None + return x + 1 # stands for this rank's shard GEMM + + +class _Original(nn.Module): + def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits=False): + return hidden_states * 0 - 5 + + +def test_the_shard_processor_returns_the_local_shard_for_selected_rows() -> None: + proc = _M3ShardDraftLogits(_Original()) + x = torch.zeros(3, 4, dtype=torch.bfloat16) + assert torch.equal(proc(x, _Head(), None, True), x + 1) + # Rows not selected by the caller take the original processor (it selects them itself). + assert torch.equal(proc(x, _Head(), None, False), x * 0 - 5) + # The shared original is not a submodule of the swap (the target model owns it). + assert list(proc.children()) == [] + + +def _worker(**over): + w = M3Eagle3Worker.__new__(M3Eagle3Worker) + if isinstance(w, nn.Module): + nn.Module.__init__( + w + ) # attribute storage only; the worker's own init needs a full model config + w.is_mtp_eagle = over.get("mtp", False) + w.guided_decoder = over.get("guided", None) + w.mapping = SimpleNamespace(enable_attention_dp=over.get("adp", False), tp_size=2, tp_rank=0) + return w + + +def _draft(head=None): + return SimpleNamespace( + lm_head=head if head is not None else _Head(), logits_processor=_Original() + ) + + +def test_the_wide_draft_logits_engage_only_where_the_shard_argmax_is_valid() -> None: + meta = SimpleNamespace(wants_advanced_draft_sampling=False) + ok = _worker()._m3_wide_draft_logits_ok(meta, _draft()) + assert ok == bool(dispatch.WIDE_SPEC_TAIL) + if not dispatch.WIDE_SPEC_TAIL: + return + assert not _worker()._m3_wide_draft_logits_ok( + SimpleNamespace(wants_advanced_draft_sampling=True), _draft() + ) + assert not _worker(guided=object())._m3_wide_draft_logits_ok(meta, _draft()) + assert not _worker(adp=True)._m3_wide_draft_logits_ok(meta, _draft()) + assert not _worker(mtp=True)._m3_wide_draft_logits_ok(meta, _draft()) + assert not _worker()._m3_wide_draft_logits_ok(meta, None) + padded = _Head() + padded.padding_size = 8 + assert not _worker()._m3_wide_draft_logits_ok(meta, _draft(padded)) + + +def test_the_base_draft_loop_routes_the_head_through_the_swapped_processor() -> None: + """The swap only matters if the in-tree loop calls draft_model.logits_processor and then sample_draft_tokens; if + the base loop stops doing so, W2b goes inert again (the frozen 2e0a73900b was: the head's processor all-gathered + and cast to fp32, so the sharded check never passed). Fail here instead of silently in-model.""" + import inspect + + from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + + # _forward_draft_loop dispatches to the linear loop, which runs the head and the pick. + dispatcher = inspect.getsource(Eagle3OneModelWorker._forward_draft_loop) + assert "self._forward_linear_draft_loop(" in dispatcher + src = inspect.getsource(Eagle3OneModelWorker._forward_linear_draft_loop) + assert "draft_model.logits_processor(" in src + assert "self.sample_draft_tokens(" in src + + +def test_a_wide_draft_pick_reaches_the_shard_argmax(monkeypatch) -> None: + """The path the in-tree loop takes on a wide step: the swapped processor returns the BF16 shard, and + sample_draft_tokens -> greedy_sample_draft_with_tp_gather hands it to the fused shard argmax, whose tokens come + back (engagement at wide rows, without a GPU).""" + if not dispatch.WIDE_SPEC_TAIL: + pytest.skip("WIDE_SPEC_TAIL is off") + calls = [] + + def fake_shard_accept(logits, vocab_offset, mapping, draft, batch_idx, force_total, *a, **k): + calls.append((tuple(logits.shape), logits.dtype, vocab_offset, tuple(draft.shape))) + m = logits.shape[0] + acc = (torch.arange(m, dtype=torch.int32) + 7).view(m, 1) + return acc, torch.ones(m, dtype=torch.int32), torch.zeros(m, dtype=torch.int64) + + monkeypatch.setattr(dispatch.m3_shard_accept_kernel, "shard_accept", fake_shard_accept) + w = _worker() + w._d2t = None + w.mapping.tp_rank = 1 + shard, rows = 100032, 46 + meta = SimpleNamespace(wants_advanced_draft_sampling=False, draft_vocab_size=2 * shard) + draft_model = _draft() + assert w._m3_wide_draft_logits_ok(meta, draft_model) + + class _ShardHead(_Head): + def apply_linear(self, x, bias): + return torch.zeros(x.shape[0], shard, dtype=torch.bfloat16) + + proc = _M3ShardDraftLogits(draft_model.logits_processor) + logits = proc(torch.zeros(rows, 8), _ShardHead(), None, True) + tokens = w.sample_draft_tokens(logits, meta, rows, draft_step=0) + assert calls == [((rows, shard), torch.bfloat16, shard, (rows, 0))] + assert torch.equal(tokens, torch.arange(rows, dtype=torch.int32) + 7) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..9e16b0a9fc2b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic layer's feed-forward exit on Eagle3 capture layers, in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``). + +Every layer's MoE / MLP output goes unreduced to the fused boundary (all-reduce + residual + the next layer's input +norm, the final norm after the last layer), and a capture layer captures the reduced residual sum that boundary +writes, with no residual of its own. CPU only: the attention, the feed-forward and the boundaries are stand-ins. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) + +LAYERS = 4 + + +class _Spec: + def __init__(self, layers): + self.layers = set(layers) + + def is_layer_capture(self, i): + return i in self.layers + + +@pytest.fixture(params=["prefill", "decode"]) +def layer(request, monkeypatch): + """(core, the target's generic layer bound to a target over ``core``).""" + calls = [] + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + c = SimpleNamespace( + w=w, + num_layers=LAYERS, + moe_layers=frozenset({2, 3}), + rms_eps=1e-6, + calls=calls, + # The 9..48-row MC finalize boundary declines on these stand-ins (the stand-in boundary is under test). + _decode=SimpleNamespace(_generic_mc_finalize=lambda i, h, r, norm_w: None), + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, + "_native_allreduce_norm", + lambda core, x, r, norm: (("normed", x, norm), ("updated", x, r)), + ) + + def boundary(self, x, x2, residual, norm_w): + calls.append((x, x2, norm_w)) + return ("next", norm_w), ("sum", x, x2) + + if request.param == "prefill": + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i))) + monkeypatch.setattr( + cls, "_swiglu_mlp", lambda self, h, prefix, low_m, dense=False: ("mlp", prefix) + ) + else: + cls = modeling.DecodeTarget + monkeypatch.setattr( + cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i), False, False) + ) + monkeypatch.setattr( + cls, + "_swiglu_mlp", + lambda self, h, prefix, low_m, push=None, dense=False: (("mlp", prefix), False), + ) + # The 33..48-row MC finalize declines here: the stand-in boundary is the one under test. + monkeypatch.setattr(cls, "_generic_mc_finalize", lambda self, i, h, r, norm_w: None) + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + monkeypatch.setattr(cls, "_generic_boundary", boundary) + target = cls(c) + return c, target._layer if request.param == "prefill" else target._generic_layer + + +@pytest.mark.parametrize("i", range(LAYERS)) +def test_a_capture_layer_captures_the_boundary_residual_sum(layer, i) -> None: + core, run = layer + captured = [] + out = run( + i, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([i]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + norm = core.w[f"l{i + 1}_in_norm"] if i + 1 < LAYERS else "final" + assert core.calls[-1][2] == norm, ( + "the boundary applies the next layer's norm (the final norm after the last)" + ) + assert captured == [(i, out[1], None)] + if i in core.moe_layers: + assert core.calls[-1][1] == ("routed", i), "the routed part folds into the boundary" + + +def test_a_layer_without_capture_captures_nothing(layer) -> None: + _, run = layer + captured = [] + run( + 1, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([]), + lambda *a: captured.append(a), + ) + assert not captured + + +def test_a_prefill_moe_layer_takes_the_mc_finalize_boundary_when_it_folds(monkeypatch) -> None: + """PrefillTarget at a MoE layer: the decode target's generic MC finalize (9..48 rows) ends the layer when it folds, + its updated residual is the Eagle3 capture, and the generic boundary does not run.""" + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + seen = [] + + def mc_finalize(i, h, r, norm_w): + seen.append((i, h, r, norm_w)) + return ("normed", i), ("updated", i) + + c = SimpleNamespace( + w=w, + num_layers=LAYERS, + moe_layers=frozenset({2, 3}), + rms_eps=1e-6, + _decode=SimpleNamespace(_generic_mc_finalize=mc_finalize), + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, "_native_allreduce_norm", lambda core, x, r, norm: (("normed", x), ("upd", x, r)) + ) + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + + def no_boundary(*a, **k): + raise AssertionError("the generic boundary runs after a folded MC finalize") + + monkeypatch.setattr(cls, "_generic_boundary", no_boundary) + monkeypatch.setattr( + cls, "_moe", lambda self, *a, **k: pytest.fail("the MoE runs outside the fold") + ) + target = cls(c) + captured = [] + out = target._layer( + 2, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([2]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + assert out == (("normed", 2), ("updated", 2)) + assert seen and seen[0][0] == 2 and seen[0][3] == "in_norm3" + assert captured == [(2, ("updated", 2), None)] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py new file mode 100644 index 000000000000..1394b8bd5924 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A dense layer's MLP runs its SwiGLU-OAI activation through ``_dense_swiglu``, the seam of the dense swigluoai C++ +op (``m3_dense_swiglu``), in the decode-graph layer (``DecodeTarget._layer``) and in the prefill layer, eager or in a +decode graph (r4.10.1 called the module directly, so the op never ran in a decode graph). The shared expert keeps the +plain ``m3_silu_and_mul``. + +CPU only: the GEMMs, the activations and the boundaries are stubbed; the test reads which seam each call takes.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, +) + + +@pytest.fixture +def seams(monkeypatch): + """Records which activation seam each MLP call takes; the GEMMs pass their input through.""" + calls = [] + monkeypatch.setattr( + modeling, + "_linear", + lambda core, x, w, s, low_m_gemm, push=None: (x, push is not None), + ) + monkeypatch.setattr( + modeling, "_dense_swiglu", lambda core, gu: calls.append(("dense", gu.shape[0])) or gu + 1 + ) + monkeypatch.setattr( + modeling, + "m3_silu_and_mul", + lambda gu, *args: calls.append(("plain", gu.shape[0])) or gu + 2, + ) + return calls + + +def _core(): + return SimpleNamespace( + swiglu_limit=7.0, + swiglu_alpha=1.702, + swiglu_beta=1.0, + num_layers=2, + moe_layers=set(), + w={"l0_post_norm": None, "l1_in_norm": None}, + ) + + +@pytest.mark.parametrize("decode_graph", [True, False]) +@pytest.mark.parametrize("rows", [1, 4, 8]) +def test_decode_dense_layer_runs_through_dense_swiglu(monkeypatch, seams, rows, decode_graph): + monkeypatch.setattr( + modeling, "flashinfer_mxfp8_decode_graph_capture_active", lambda: decode_graph + ) + target = SimpleNamespace(core=_core()) + target._attention = lambda i, position_ids, hidden, step, armed: (hidden, False) + target._attn_push_target = lambda num_tokens: None + target._boundary = lambda mc_ar, h, residual, norm, **kwargs: (h, residual) + target._swiglu_mlp = lambda *args, **kwargs: modeling.DecodeTarget._swiglu_mlp( + target, *args, **kwargs + ) + x = torch.zeros(rows, 16) + plan = SimpleNamespace(attn_mc_ar=False, moe_mc_ar=False, qkv_prequant=False) + out, _ = modeling.DecodeTarget._layer(target, 0, plan, None, x, x.clone(), {}) + assert seams == [("dense", rows)] + assert torch.equal(out, x + 1) + + +@pytest.mark.parametrize("rows", [1, 64]) +def test_prefill_dense_mlp_runs_through_dense_swiglu(seams, rows): + target = SimpleNamespace(core=_core()) + x = torch.zeros(rows, 16) + out = modeling.PrefillTarget._swiglu_mlp(target, x, "l0_mlp", True, dense=True) + assert seams == [("dense", rows)] + assert torch.equal(out, x + 1) + + +def test_shared_expert_keeps_the_plain_activation(seams): + target = SimpleNamespace(core=_core()) + x = torch.zeros(4, 16) + out, pushed = modeling.DecodeTarget._swiglu_mlp(target, x, "l3_sh", True) + assert seams == [("plain", 4)] and not pushed + assert torch.equal(out, x + 2) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..c9cd78d40985 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,771 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp2 (Rubin) decode target's fast-path table (``modeling.py``: the row ranges and the selection functions): its +row ranges against the catalog cells that certify them, and the selection for every decode row count. + +Every range's covered shapes are listed in ``ENABLED``; a range not listed covers nothing. + +CPU only: nothing here launches a kernel. +""" + +import inspect +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_draft_step as draft_step_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + triton_sparse_decode as tsd, +) + +ROWS = range(1, dispatch.TINY_MAX_TOKENS + 2) # every decode row count and the first one past it + + +@pytest.fixture(autouse=True) +def _policy(): + """Every test starts outside require mode with nothing logged, and leaves it so.""" + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_every_covered_shape_is_a_certified_cell(fast_path) -> None: + uncertified = fast_path.keys - fast_path.certified() + assert not uncertified, ( + f"{fast_path.path} covers {sorted(uncertified, key=str)} but " + f"{type(fast_path.entry).__name__}.CELLS certifies none of them: add a cell or narrow the range" + ) + + +def test_the_fast_path_names_are_unique() -> None: + names = [r.path for r in dispatch.FAST_PATHS] + assert len(names) == len(set(names)) + + +def test_the_decode_target_takes_the_gen_workers_batch() -> None: + """Up to eight requests verifying 1 + 3 drafts (32 rows) take the decode target; a wider step is the prefill + target's. Every fast path still keeps its own rows (the batch-1/2 target's TINY_MAX_TOKENS unless raised).""" + assert dispatch.TINY_MAX_TOKENS == 8 + rows = range(1, dispatch.DECODE_MAX_TOKENS + 2) + if dispatch.B8_DECODE: + assert dispatch.DECODE_MAX_TOKENS == 32 + elif dispatch.B4_DECODE: + assert dispatch.DECODE_MAX_TOKENS == 16 + else: + assert dispatch.DECODE_MAX_TOKENS == dispatch.TINY_MAX_TOKENS + expected = [True] * dispatch.DECODE_MAX_TOKENS + [False] + assert [dispatch.is_decode_step(n) for n in rows] == expected + + +# ------------------------------------------------------------------ the cells turned on + + +def test_the_moe_engine_takes_every_decode_step() -> None: + """The MoE lane's m3_moe engine (TP2 width 1536, shared expert folded in) at 1..8 rows, unless killed; its sum + pushes into the multicast boundary all-reduce at the engine's rows while that all-reduce is on.""" + expected = set() if dispatch.KILL_MOE_ENGINE else set(range(1, 9)) + assert set(dispatch.MOE_ENGINE.keys) == expected + for rows in ROWS: + if rows in expected: + want = dispatch.MOE_ENGINE.path + elif dispatch.MC_FINALIZE.covers(rows): # B3: past the engine, the multicast finalize + want = dispatch.MC_FINALIZE.path + else: + want = None + assert dispatch.moe_path(rows) == want, rows + push_rows = dispatch.MOE_ENGINE_MAX_TOKENS if dispatch.MC_AR_TP2 else 0 + assert dispatch.MOE_PUSH_MAX_TOKENS == push_rows + assert set(dispatch.MOE_PUSH.keys) == set(range(1, push_rows + 1)) + + +# ------------------------------------------------------------------ on: the sparse layers' block selection + + +def test_the_index_scorer_takes_one_to_forty_eight_requests_with_their_schedules() -> None: + # W-IDX1 (9..48 requests) is parked: the default covers 1..8; its schedules below stay certified for the knob. + assert dispatch.INDEX_SCORE.keys == frozenset(range(1, 9)) + assert dispatch.index_score_schedule(1) == (1536, True) + assert dispatch.index_score_schedule(2) == (768, True) + assert dispatch.index_score_schedule(3) == (424, True) + assert dispatch.index_score_schedule(4) == (318, True) + for requests in range(5, 9): + assert dispatch.index_score_schedule(requests) == (1272 // requests, True) + # W-IDX1: 9..48 requests in 7 split buckets, about 1272 CTAs in total, early issue + assert len({dispatch.index_score_schedule(r)[0] for r in range(9, 49)}) == 7 + for requests in range(9, 49): + split, early = dispatch.index_score_schedule(requests) + assert early and requests * split <= 1272 + split + # every covered request count has a split (the graph captures it) + assert set(dispatch.INDEX_SCORE.keys) <= set(dispatch.INDEX_SCORE_SPLIT_K) + + +def test_the_long_select_takes_every_key_width_up_to_its_maximum() -> None: + assert dispatch.LONG_SELECT.keys == frozenset(dispatch.LONG_SELECT_KEY_WIDTHS) + assert (dispatch.LONG_SELECT_MIN_BLOCKS, dispatch.LONG_SELECT_MAX_BLOCKS) == (129, 9216) + + +def test_the_384_thread_select_takes_every_key_width_with_its_launch_key() -> None: + assert dispatch.SELECT_LONG_DSL is True + assert dispatch.LONG_SELECT_DSL.keys == frozenset(dispatch.LONG_SELECT_DSL_KEY_WIDTHS) + keys = dispatch.long_select_dsl_keys() + assert [key[1:] for key in keys] == [(False, dispatch.SELECT_LONG_TRIGGER_PHASE)] + + +@pytest.mark.parametrize("num_blocks", [129, 1294, 5888, 9216]) +def test_a_long_select_decline_inside_its_range_is_a_miss(num_blocks) -> None: + """The kernel declining operands its cells certify must not fall back silently under require.""" + assert dispatch.long_select_covers(num_blocks) + # Outside require: logged once, the in-tree selector runs. + dispatch.long_select_fallback(num_blocks, 8) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="long_select"): + dispatch.long_select_fallback(num_blocks, 8) + + +@pytest.mark.parametrize("num_blocks", [1, 16, 128, 9217, 20000]) +def test_rows_outside_the_long_select_range_are_uncovered_even_under_require(num_blocks) -> None: + """Short rows and rows past LONG_SELECT_MAX_BLOCKS take the in-tree selector in every mode (long contexts must not + fail under require).""" + assert not dispatch.long_select_covers(num_blocks) + dispatch.set_require(True) + dispatch.long_select_fallback(num_blocks, 8) + + +def test_taken_logs_without_raising_under_require() -> None: + dispatch.set_require(True) + dispatch.taken(dispatch.INDEX_SCORE.path, 1) + dispatch.taken(dispatch.LONG_SELECT.path, 8) + + +def test_an_index_score_cell_of_another_schedule_certifies_nothing() -> None: + spec = dict(seq_lens=(165521,), qlen=4, heads=2, split_k=512, early=True) + assert tuple(dispatch.INDEX_SCORE.certifies(spec)) == () + assert tuple(dispatch.INDEX_SCORE.certifies(dict(spec, split_k=1536, heads=1))) == () + assert tuple(dispatch.INDEX_SCORE.certifies(dict(spec, split_k=1536))) == (1,) + + +#: Every TP2 MXFP8 projection (N, K): sparse-layer QKV + index, dense QKV, o_proj, shared gate_up / down, dense +#: gate_up / down. +TP2_PROJECTIONS = ( + (4992, 6144), + (4608, 6144), + (6144, 4096), + (3072, 6144), + (6144, 1536), + (12288, 6144), + (6144, 6144), +) +_DECODE_ROWS = range(1, dispatch.TINY_MAX_TOKENS + 1) +# Requests of a fused spec-tail step (B4 spec tail: up to the target's steps). +_SPEC_REQS = range(1, dispatch.SPEC_TAIL_MAX_REQUESTS + 1) +# W2: the fused verify's requests (up to WIDE_MAX_TOKENS rows, unless WIDE_SPEC_TAIL is killed). +_VERIFY_REQS = range(1, dispatch.WIDE_VERIFY_MAX_REQUESTS + 1) +#: B2: the GEMVs' rows at batch 4 (1..16 unless MX_GEMV_B4 is killed). The bf16 GEMV declines dense gate_up / down at +#: 9..16 rows (24 k-tiles per split at MMA width 16, past its in-kernel quantize limit of 16). +_GEMV_HI = ( + (dispatch.DECODE_MAX_TOKENS if dispatch.B8_MX_GEMV_32 else dispatch.B4_MAX_TOKENS) + if dispatch.MX_GEMV_B4 + else dispatch.TINY_MAX_TOKENS +) +_BF16_WIDE_DECLINES = frozenset({(12288, 6144), (6144, 6144)}) + + +#: Batch 8: at 17..32 rows (MMA width 32) the bf16 GEMV's in-kernel quantize serves only shared down (6 k-tiles per +#: split); the other projections keep their 16-row reach. +_BF16_B8_SERVES = frozenset({(6144, 1536)}) + + +#: Step B (GEMM-AR 794929): dense gate_up / down and o_proj keep the 16-row GEMV even with B8_MX_GEMV_32. +_GEMV_B4_ROWS_ONLY = frozenset({(12288, 6144), (6144, 6144), (6144, 4096)}) + + +def _mx_hi(nk) -> int: + if nk in _GEMV_B4_ROWS_ONLY: + return min(_GEMV_HI, 16) + # W-G1: the four projections it raises reach their served rows up to 48 (dispatch._GEMV_W48_ROWS). + return max(dispatch._GEMV_W48_ROWS.get(nk, ()), default=_GEMV_HI) + + +def _bf16_hi(nk) -> int: + if nk in _BF16_WIDE_DECLINES: + return dispatch.TINY_MAX_TOKENS + if _GEMV_HI > 16 and nk not in _BF16_B8_SERVES: + return 16 + return _GEMV_HI + + +#: Fast paths a lane has turned on, each with the shapes it covers (the rest of the skeleton stays off). +ENABLED = { + # GEMM-AR (agg G-d): the tensor-core router of an unfused MoE at 9..48 rows, as the agg canon runs it. + "m3_router_tc": frozenset( + range(dispatch.ROUTER_TC_GEN_MIN_ROWS, dispatch.ROUTER_TC_GEN_MAX_ROWS + 1) + ), + # SpecTail: sampled (temperature + top-p) verify of 1-2 requests at draft length 3 (TP2 shard kernel). + "m3_sampled_accept": frozenset( + (3, n) for n in range(1, dispatch.SAMPLED_VERIFY_MAX_REQUESTS + 1) + ), + # SpecTail: greedy verify (shard argmax + accept) and the fused draft loop of 1-2 requests at draft length 3, + # and that loop's draft-input rows (1 / 2 later-pass rows, 4 / 8 first-pass rows). + "m3_shard_accept": frozenset((3, n) for n in _VERIFY_REQS), + "m3_draft_step": frozenset((3, n) for n in _SPEC_REQS), + "m3_draft_input": frozenset(r for n in _SPEC_REQS for r in (n, 4 * n)), + # SpecTail: the drafter's BF16 linears on FlashInfer's low-M kernels (TP8 port) at the certified b1 rows. + "m3_draft_lowm": frozenset( + { + ((4608, 12288), 1), + ((4608, 12288), 4), + ((6144, 4096), 1), + ((6144, 4096), 4), + ((18432, 6144), 1), + ((6144, 9216), 1), + ((6144, 9216), 4), + ((6144, 18432), 1), + } + # Batch 4: the 16-row draft pass 0 (o, down, fc), unless DRAFT_LOWM_MAX_ROWS is set back to 4. + | ( + {((6144, 4096), 16), ((6144, 9216), 16), ((6144, 18432), 16)} + if dispatch.DRAFT_LOWM_MAX_ROWS >= 16 + else set() + ) + # Batch 8: the 8-row passes' o and the 32-row pass 0's down, unless DRAFT_LOWM_MAX_ROWS is set lower. + | ({((6144, 4096), 8)} if dispatch.DRAFT_LOWM_MAX_ROWS >= 8 else set()) + | ({((6144, 9216), 32)} if dispatch.DRAFT_LOWM_MAX_ROWS >= 32 else set()) + ), + # SpecTail: the fused first drafter inputs of 1-2 fused-verify requests, the capture norm and the chunk norm at + # their 4 / 8 verify rows (own ranges and kills since the TP2 fast-path audit). + "m3_fused_first_inputs": frozenset(_SPEC_REQS), + "m3_capture_norm": frozenset(4 * n for n in _SPEC_REQS), + "m3_chunk_norm": frozenset(4 * n for n in _SPEC_REQS), + # MSA: the long-context sparse decode with the MXFP8 o_proj hand-off, every decode step, and the sparse layers' + # NVFP4 QKV + indexer producer at the TP2 heads. + "m3_sparse_decode_lc": frozenset({"mxfp8"}), + "m3_producer": frozenset({"nvfp4_tp2"}), + "sparse_decode_producer": frozenset({"per_row"}), + # MoE: the m3_moe engine (TP2 width 1536, shared expert folded in) at 1..8 rows, unless killed. + "m3_moe_engine": frozenset() if dispatch.KILL_MOE_ENGINE else frozenset(range(1, 9)), + # Indexer: the index scorer at 1-8 requests, the long-row select at every key width (129..9216 blocks). + "m3_index_score": frozenset(range(1, dispatch.INDEX_SCORE_MAX_REQUESTS + 1)), + "long_select": frozenset({2, 4, 9}), + "long_select_dsl": frozenset({4, 11, 24}), + # LongAttn: the dense layers' full-context decode of 1-2 requests, 1 + 3 verify rows or one draft-pass row. + "m3_fullctx_decode": frozenset({(4, 1), (1, 1), (4, 2), (1, 2)}), + # MoE (uGPU): the die-placed target head at 4 / 8 / 12 / 16 verify rows, unless killed. + "m3_ugpu_head": frozenset() if dispatch.KILL_UGPU_HEAD else frozenset({4, 8, 12, 16}), + # GEMM/AR (+ B2): the decode MXFP8 GEMVs at every TP2 projection and decode row, on an MXFP8 and a bf16 + # activation, raised to 16 rows where the kernel serves them. + "m3_mx_gemv": frozenset((nk, m) for nk in TP2_PROJECTIONS for m in range(1, _mx_hi(nk) + 1)) + if dispatch.MX_GEMV_TP2 + else frozenset(), + "m3_mx_gemv_bf16": frozenset( + (nk, m) for nk in TP2_PROJECTIONS for m in range(1, _bf16_hi(nk) + 1) + ) + if dispatch.MX_GEMV_TP2 + else frozenset(), + # GEMM/AR: the multicast boundary all-reduce, the o_proj's push into it, the fused embedding + first RMSNorm. + # (+ B3: 9..16 rows on the 16-token bucket, unless B4_MC_AR is killed.) + "mc_allreduce_norm": frozenset(range(1, dispatch.PLAN_MC_AR_MAX_TOKENS + 1)) + if dispatch.MC_AR_TP2 + else frozenset(), + # (+ B6: 9..16 rows into the 16-token bucket, unless B4_ATTN_PUSH is killed; the o_proj GEMV stops at 16.) + "attn_push": frozenset( + range(1, min(dispatch.PLAN_MC_AR_MAX_TOKENS, dispatch.B4_MAX_TOKENS) + 1) + if dispatch.B4_ATTN_PUSH + else _DECODE_ROWS + ) + if (dispatch.MC_AR_TP2 and dispatch.ATTN_PUSH_TP2) + else frozenset(), + "dense_push": frozenset(_DECODE_ROWS) + if (dispatch.MC_AR_TP2 and dispatch.DENSE_PUSH_TP2) + else frozenset(), + # B3: the MoE's multicast finalize + all-reduce + norm at 9..16 rows (past the engine), unless killed. + "mc_finalize_allreduce_norm": frozenset( + range(dispatch.MC_FINALIZE_MIN_TOKENS, dispatch.PLAN_MC_AR_MAX_TOKENS + 1) + ) + if dispatch.B4_MC_FINALIZE + else frozenset(), + # MoE: the engine's push into the multicast boundary all-reduce, at the engine's rows. + "moe_push": frozenset(range(1, dispatch.MOE_PUSH_MAX_TOKENS + 1)), + # (+ B5: 9..16 rows, unless B4_EMBED_FIRSTNORM is killed; 17..32 keep the stock layer 0.) + "embed_firstnorm": frozenset( + range(1, dispatch.B4_MAX_TOKENS + 1) if dispatch.B4_EMBED_FIRSTNORM else _DECODE_ROWS + ) + if dispatch.EMBED_FIRSTNORM_TP2 + else frozenset(), + # SpecTail: the fused draft loop's all-reduce boundaries (1..8 rows) on the multicast all-reduce (opts 16). + "m3_draft_mc_allreduce_norm": frozenset(range(1, dispatch.DRAFT_MC_AR_MAX_ROWS + 1)), + # LongAttn: the Eagle3 drafter layer's generation passes of 1-2 requests on the same recipes. + "m3_drafter_fullctx": frozenset({(4, 1), (1, 1), (4, 2), (1, 2)}), +} + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_skeleton_no_fast_path_covers_a_shape(fast_path) -> None: + expected = ENABLED.get(fast_path.path, frozenset()) + assert fast_path.keys == expected, f"{fast_path.path} covers {sorted(fast_path.keys, key=str)}" + + +def test_the_sparse_decode_writes_the_mxfp8_handoff() -> None: + """The sparse decode (LC2, the Triton merge) writes the MXFP8 o_proj hand-off; the range keys the target's + SPARSE_DECODE_LC / SPARSE_DECODE_PRODUCER cover are the kernel module's.""" + assert tsd.MERGE_MXFP8_HANDOFF is True + assert dispatch.SPARSE_DECODE_LC.keys is tsd.LC_KEYS + assert dispatch.SPARSE_DECODE_PRODUCER.keys is tsd.PRODUCER_KEYS + + +def test_sparse_decode_lc_covers_the_decode_steps_with_the_handoff() -> None: + """Every decode step (at most DECODE_MAX_TOKENS rows) writes the MXFP8 o_proj hand-off, so the long-context + sparse decode needs only its "mxfp8" variant, on the stock log-domain recurrence and without pre-wait staging.""" + assert dispatch.SPARSE_DECODE_LC.keys == frozenset({"mxfp8"}) + assert tsd.MERGE_MXFP8_MIN_TOKENS <= 1 and dispatch.DECODE_MAX_TOKENS <= tsd.MergeMxfp8.MAX_ROWS + assert tsd.LC_PREWAIT is False + + +def test_the_dense_layers_take_the_fullctx_decode_at_the_gen_workers_shapes() -> None: + """One or two requests, 1 + 3 verify rows or one draft-pass row, each at its own pinned recipe.""" + assert dispatch.FULLCTX_DECODE_ENABLED + assert dispatch.FULLCTX_DECODE.keys == {(4, 1), (1, 1), (4, 2), (1, 2)} + assert set(dispatch.FULLCTX_DECODE_RECIPES) == set(dispatch.FULLCTX_DECODE.keys) + for recipe in dispatch.FULLCTX_DECODE_RECIPES.values(): + assert 2 <= recipe["splits"] <= 128 # the standalone reducer's range + + +def test_the_fullctx_entry_cells_run_at_the_forward_plan_capacity() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention import ( # noqa: E501 + m3_fullctx_decode as entry, + ) + + assert entry.MAX_KV_LEN == dispatch.FULLCTX_DECODE_MAX_KV_LEN + + +def test_the_fullctx_capacity_covers_the_agentx_cache_bound() -> None: + """The forward compares the cache's length bound (max_seq_len plus the speculative rows it reserves) with the plan + capacity; at the TP2 AgentX max_seq_len of 1 << 20 with up to 8 draft rows the dense layers must still be taken.""" + assert dispatch.FULLCTX_DECODE_MAX_KV_LEN >= (1 << 20) + 8 + + +def test_the_drafter_takes_the_fullctx_decode_at_the_same_shapes() -> None: + assert dispatch.DRAFTER_FULLCTX_ENABLED + assert dispatch.DRAFTER_FULLCTX.keys == dispatch.FULLCTX_DECODE.keys + + +@pytest.mark.parametrize("rows", range(1, 50)) +def test_the_tensor_core_router_takes_9_to_48_rows(rows) -> None: + """agg G-d: the tensor-core router of an unfused MoE at 9..48 rows, as the agg canon runs it; 1..8 rows stay on + the MoE engine's fused router (the batch-1/2 target's) and 49+ on the fp32 router.""" + assert dispatch.ROUTER_TC_GEN.covers(rows) == (9 <= rows <= 48) + + +# ------------------------------------------------------------------ landed: the decode MXFP8 GEMVs (GEMM/AR lane) + + +def test_the_gemv_rows_follow_the_kill_constant() -> None: + expected = set(TP2_PROJECTIONS) if dispatch.MX_GEMV_TP2 else set() + assert set(dispatch.MX_GEMV_ROWS) == expected + assert set(dispatch.MX_GEMV_BF16_MAX_TOKENS) == expected + + +@pytest.mark.parametrize("rows", range(1, dispatch.TINY_MAX_TOKENS + 1)) +@pytest.mark.parametrize("nk", TP2_PROJECTIONS, ids=str) +def test_every_projection_takes_the_gemv_at_every_decode_row(nk, rows) -> None: + """MXFP8 activation (a producer quantized it) -> the GEMV; bf16 -> the GEMV quantizing in its prologue (staged: + the stock decode-graph GEMMs while MX_GEMV_TP2 is off).""" + on = dispatch.MX_GEMV_TP2 + assert dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) == ( + dispatch.MX_GEMV_PATH if on else dispatch.PREQUANT_GEMM_PATH + ) + assert dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) == ( + dispatch.MX_GEMV_BF16_PATH if on else dispatch.GENERIC_GEMM_PATH + ) + + +@pytest.mark.parametrize("nk", TP2_PROJECTIONS, ids=str) +def test_wider_than_the_decode_step_is_the_generic_gemm(nk) -> None: + # Past the decode step and past every GEMV row it serves (W-G1 raises four projections to 48). + rows = max(dispatch.DECODE_MAX_TOKENS, _mx_hi(nk)) + 1 + assert ( + dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) + == dispatch.PREQUANT_GEMM_PATH + ) + assert ( + dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) + == dispatch.GENERIC_GEMM_PATH + ) + + +def test_an_unknown_projection_is_the_generic_gemm() -> None: + assert ( + dispatch.linear_path(256, 512, 1, prequant=True, low_m_gemm=True) + == dispatch.PREQUANT_GEMM_PATH + ) + assert ( + dispatch.linear_path(256, 512, 1, prequant=False, low_m_gemm=True) + == dispatch.GENERIC_GEMM_PATH + ) + + +def test_taken_logs_once_per_path_and_rows(monkeypatch) -> None: + lines = [] + monkeypatch.setattr(dispatch.logger, "info", lines.append) + for _ in range(3): + dispatch.taken(dispatch.MX_GEMV_PATH, 4, "[6144, 4096] mxfp8 in") + dispatch.taken(dispatch.MX_GEMV_PATH, 8, "[6144, 4096] mxfp8 in") + assert len(lines) == 2 and all("m3_mx_gemv taken at" in line for line in lines) + + +def test_skeleton_the_dense_decode_split_is_the_built_in_one() -> None: + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {} + + +def test_skeleton_speculation_runs_the_in_tree_loop() -> None: + # SpecTail: fused greedy / sampled verify and the fused draft loop up to the target's steps (B4 spec tail; the + # batch-1/2 rows with B4_SPEC_TAIL off; BF16 head; no coded heads). + reqs = dispatch.SPEC_TAIL_MAX_TOKENS // (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert dispatch.SPEC_TAIL_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if dispatch.B4_SPEC_TAIL else dispatch.TINY_MAX_TOKENS + ) + # W2: the fused verify alone widens to WIDE_MAX_TOKENS rows (WIDE_SPEC_TAIL off keeps SPEC_TAIL_MAX_TOKENS). + verify = ( + dispatch.WIDE_MAX_TOKENS if dispatch.WIDE_SPEC_TAIL else dispatch.SPEC_TAIL_MAX_TOKENS + ) // (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert dispatch.WIDE_VERIFY_MAX_REQUESTS == verify + assert dispatch.SHARD_ACCEPT_MAX_REQUESTS == verify + # The sampled verify stays at SPEC_TAIL_MAX_TOKENS rows (past it, the in-tree sampler). + assert dispatch.SAMPLED_VERIFY_MAX_REQUESTS == reqs + assert dispatch.DRAFT_STEP_MAX_REQUESTS == min(reqs, 8) + assert dispatch.FC_EARLY_MAX_ROWS == 0 + assert not draft_step_kernel.HEAD_CODED + + +def test_a_miss_raises_only_under_require() -> None: + dispatch.missed("some_path", 1, "a precondition") + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="some_path"): + dispatch.missed("some_path", 1, "a precondition") + + +def test_skeleton_first_drafter_inputs_are_in_tree() -> None: + # SpecTail turned the sampled verify on (SAMPLED_VERIFY_MAX_REQUESTS 2, test_skeleton_speculation_...). + assert dispatch.FIRST_DRAFTER_INPUTS_MAX_REQUESTS == 0 + + +@pytest.mark.parametrize( + "rows, capture_normed, taken", + [ + (4, False, True), + (8, False, True), + (8, True, True), + # 3 / 4 requests: the B4 spec tail's rows (in-tree with B4_SPEC_TAIL off). + (12, False, dispatch.B4_SPEC_TAIL), + (12, True, dispatch.B4_SPEC_TAIL), + (16, True, dispatch.B4_SPEC_TAIL), + ], +) +def test_the_fused_first_inputs_take_only_their_covered_rows(rows, capture_normed, taken) -> None: + """The fused first drafter inputs (and the fc fork) take a fused-verify step only at the requests + FUSED_FIRST_INPUTS covers and, unless the captures applied the fc_norm, the rows CHUNK_NORM covers; anything else + runs the in-tree first drafter inputs.""" + + worker = object.__new__(dispatch.M3Eagle3Worker) + worker._m3_capture_norm = (None, 1e-6) if capture_normed else None + meta = SimpleNamespace(runtime_draft_len=dispatch.SHARD_ACCEPT_DRAFT_LEN) + assert worker._m3_first_inputs_covered(rows, meta) is taken + + +# ------------------------------------------------------- landed: the multicast boundary all-reduce (GEMM/AR lane) + + +def test_both_boundaries_take_the_multicast_all_reduce() -> None: + """One 8-token bucket for the whole decode step; the MoE boundary publishes the next QKV's MXFP8 rows (killed: + every boundary on the built-in all-reduce while MC_AR_TP2 is off).""" + on = dispatch.MC_AR_TP2 + tiny = dispatch.DECODE_PLANS[0] + assert (tiny.name, tiny.max_tokens) == ("tiny", dispatch.TINY_MAX_TOKENS) + assert (tiny.attn_mc_ar, tiny.moe_mc_ar, tiny.qkv_prequant) == (on, on, on), tiny + # B6: the o_proj push follows the boundary's rows up to the o_proj GEMV's 16, unless B4_ATTN_PUSH is killed. + push_rows = ( + min(dispatch.PLAN_MC_AR_MAX_TOKENS, dispatch.B4_MAX_TOKENS) + if dispatch.B4_ATTN_PUSH + else dispatch.TINY_MAX_TOKENS + ) + assert dispatch.ATTN_PUSH_MAX_TOKENS == (push_rows if on and dispatch.ATTN_PUSH_TP2 else 0) + assert not dispatch.B4_ATTN_PUSH or dispatch.B4_MC_AR + assert dispatch.ATTN_PUSH_MAX_TOKENS <= max( + r for (nk, r) in dispatch.MX_GEMV.keys if nk == dispatch.O_PROJ + ) + assert dispatch.DENSE_PUSH_MAX_TOKENS == ( + dispatch.TINY_MAX_TOKENS if on and dispatch.DENSE_PUSH_TP2 else 0 + ) + for rows in _DECODE_ROWS: + assert dispatch.plan_for(rows) is dispatch.DECODE_PLANS[0] + # The MoE pushes only from its engine, at the engine's rows. + assert dispatch.MOE_PUSH_MAX_TOKENS == (dispatch.MOE_ENGINE_MAX_TOKENS if on else 0) + + +def test_the_generic_boundaries_take_the_multicast_all_reduce() -> None: + """Prefill, mixed steps and the Eagle3 capture layers: the multicast all-reduce's 48-token bucket, alone and folded, + and the MoE finalize folded into it, as the agg canon runs them.""" + rows = 48 if dispatch.MC_AR_TP2 else 0 + assert dispatch.GENERIC_MC_AR_MAX_TOKENS == rows + assert dispatch.GENERIC_MC_AR_FOLDED_MAX_TOKENS == rows + assert dispatch.GENERIC_MC_FINALIZE_MAX_TOKENS == rows + + +def test_the_tensor_core_router_reaches_48_rows() -> None: + assert dispatch.ROUTER_TC_GEN_MAX_ROWS == 48 + + +def test_the_sampled_verify_owner_walk_is_one_token_by_default() -> None: + assert dispatch.m3_sampled_accept_kernel.ILP_WALK is False + + +def test_the_fused_embedding_rows_follow_the_kill_constant() -> None: + # B5: 9..16 rows too (not 17..32), unless B4_EMBED_FIRSTNORM is killed. + rows = dispatch.B4_MAX_TOKENS if dispatch.B4_EMBED_FIRSTNORM else dispatch.TINY_MAX_TOKENS + assert dispatch.EMBED_FIRSTNORM_MAX_TOKENS == (rows if dispatch.EMBED_FIRSTNORM_TP2 else 0) + + +# ------------------------------------------------------- landed: MX-TAC (pinned decode-graph MXFP8 GEMM tactic) + + +def test_the_decode_mxfp8_decision_is_pinned_to_cute_dsl() -> None: + from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, + ) + + assert dispatch.MXFP8_DECODE_QUANT_TACTIC == MXFP8QuantizeRunner.CUTE_DSL + assert dispatch.MXFP8_DECODE_GEMM_TACTIC == FlashInferMXFP8GemmRunner.CUTE_DSL + + +@pytest.mark.parametrize("target", ["DecodeTarget", "PrefillTarget"]) +def test_the_moe_forward_reports_the_engine_taken(target) -> None: + """The engagement evidence comes from the forward (not the engine build's warm call).""" + + src = inspect.getsource(getattr(dispatch, target)._moe) + assert src.index("taken(MOE_ENGINE.path") < src.index("m3_moe_engine(h.contiguous()") + + +# ------------------------------------------------------------------ the drafter's low-M linears (TP8 port) + + +def test_draft_lowm_covers_exactly_the_certified_rows_of_its_tactic_table() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm import ( # noqa: E501 + m3_draft_lowm as entry, + ) + + assert set(entry.TACTICS) == set(entry.ROWS) == set(dispatch.DRAFT_LOWM_ROWS) + assert dispatch.DRAFT_LOWM.keys == frozenset( + (nk, m) + for nk, rows in entry.ROWS.items() + for m in rows + if m <= dispatch.DRAFT_LOWM_MAX_ROWS + ) + for nk, table in entry.TACTICS.items(): + rows = [row[0] for row in table] + assert rows == sorted(rows) and len(set(rows)) == len(rows), nk + for m in entry.ROWS[nk]: + assert entry.tactic_for(*nk, m) is not None, (nk, m) + + +def test_prebuild_draft_head_is_idempotent_and_keeps_replaced_copies(monkeypatch) -> None: + """A repeated cache_derived_state with an unchanged head weight keeps the built (placed) copy; a changed weight + rebuilds it and keeps the old copy alive, so a CUDA graph captured earlier never reads unmapped memory.""" + + import torch + + built = [] + + def fake_prepare(lm_head, tp_rank): + head = SimpleNamespace(weight=lm_head.weight) + built.append(head) + return head + + monkeypatch.setattr(draft_step_kernel, "prepare_head", fake_prepare) + worker = object.__new__(dispatch.M3Eagle3Worker) + worker._m3_draft_head = None + worker._m3_draft_head_version = None + worker._m3_draft_head_graveyard = [] + worker.mapping = SimpleNamespace(tp_rank=0) + monkeypatch.setattr(worker, "_m3_drafter_supported", lambda dm: True, raising=False) + model = SimpleNamespace(lm_head=SimpleNamespace(weight=torch.zeros(4, 4))) + + assert worker.prebuild_draft_head(model) and len(built) == 1 + first = worker._m3_draft_head + assert worker.prebuild_draft_head(model) and len(built) == 1 and worker._m3_draft_head is first + model.lm_head.weight.add_(1.0) # new values: version bump + assert worker.prebuild_draft_head(model) and len(built) == 2 + assert worker._m3_draft_head is built[1] and worker._m3_draft_head_graveyard == [first] + + +# ----------------------------------------------------------------------- batch 4 (this target): the b16 plan + + +def test_the_batch_8_steps_take_the_b16_and_b32_plans() -> None: + """9..16-row steps run the "b16" plan (16-token bucket) and 17..32-row steps the "b32" plan (48-token bucket), both + on the multicast all-reduce (B4_MC_AR) with no MXFP8 copy for the next QKV; 1..8 rows keep the "tiny" plan.""" + if not dispatch.B4_DECODE: + assert len(dispatch.DECODE_PLANS) == 1 + return + names = [p.name for p in dispatch.DECODE_PLANS] + assert names == (["tiny", "b16", "b32"] if dispatch.B8_DECODE else ["tiny", "b16"]) + b16 = dispatch.DECODE_PLANS[1] + assert b16.max_tokens == dispatch.B4_MAX_TOKENS == 16 + assert dispatch.DECODE_MAX_TOKENS == (32 if dispatch.B8_DECODE else 16) + b3 = dispatch.B4_MC_AR + for plan in dispatch.DECODE_PLANS[1:]: + quant = dispatch.B8_QKV_PREQUANT if plan.name == "b32" else dispatch.B4_QKV_PREQUANT + assert (plan.attn_mc_ar, plan.moe_mc_ar, plan.qkv_prequant) == (b3, b3, quant) + assert dispatch.PLAN_MC_AR_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if b3 else dispatch.TINY_MAX_TOKENS + ) + for rows in range(1, dispatch.DECODE_MAX_TOKENS + 1): + want = "tiny" if rows <= dispatch.TINY_MAX_TOKENS else "b16" if rows <= 16 else "b32" + assert dispatch.plan_for(rows).name == want, rows + + +def test_the_batch_4_gemv_rows() -> None: + """B2: the decode GEMVs serve 1..16 rows (kill MX_GEMV_B4: 1..8), the bf16 one except dense gate_up / down.""" + if not dispatch.MX_GEMV_TP2: + return + for nk, rows in dispatch.MX_GEMV_ROWS.items(): + assert rows == frozenset(range(1, _mx_hi(nk) + 1)), nk + assert dispatch.MX_GEMV_BF16_MAX_TOKENS[nk] == _bf16_hi(nk), nk + + +@pytest.mark.parametrize("nk", sorted(_BF16_WIDE_DECLINES), ids=str) +def test_the_bf16_gemv_declines_take_the_generic_gemm_at_batch_4(nk) -> None: + """Dense gate_up / down on a bf16 activation: the GEMV at 1..8 rows, the generic GEMM at 9..16 (the MXFP8 GEMV + still serves them on a prequantized activation).""" + if not dispatch.MX_GEMV_TP2: + return + for rows in range(1, dispatch.DECODE_MAX_TOKENS + 1): + bf16 = dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=False) + mx = dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=False) + assert bf16 == ( + dispatch.MX_GEMV_BF16_PATH + if rows <= dispatch.TINY_MAX_TOKENS + else dispatch.GENERIC_GEMM_PATH + ), rows + assert mx == ( + dispatch.MX_GEMV_PATH if rows <= _mx_hi(nk) else dispatch.PREQUANT_GEMM_PATH + ), rows + + +def test_no_other_fast_path_is_raised_past_the_batch_1_2_rows() -> None: + """Only the paths this target raises for 16 rows (the GEMVs, B3's multicast all-reduce) cover more than + TINY_MAX_TOKENS rows; every other row-keyed path declines a 9..16-row step to stock inside the decode target.""" + raised = { + dispatch.MX_GEMV.path, + dispatch.MX_GEMV_BF16.path, + dispatch.MC_ALLREDUCE.path, + dispatch.MC_FINALIZE.path, + dispatch.ROUTER_TC_GEN.path, # agg G-d: 9..48 rows + } + # SpecTail TR-5b: B5's fused embedding and B6's o_proj push past 8 rows, unless killed. + if dispatch.B4_EMBED_FIRSTNORM: + raised.add(dispatch.EMBED_FIRSTNORM.path) + if dispatch.B4_ATTN_PUSH: + raised.add(dispatch.ATTN_PUSH.path) + if dispatch.B4_SPEC_TAIL: + # The B4 spec tail: the fused verify's and draft loop's row-keyed paths up to SPEC_TAIL_MAX_TOKENS. + raised |= { + dispatch.DRAFT_INPUT.path, + dispatch.CAPTURE_NORM.path, + dispatch.CHUNK_NORM.path, + } + # Keys that are not row counts (both selects' key widths; the index scorer's request counts, W-IDX1), or rows + # already certified past 8 in the batch-1/2 target (the die-split head's 12 / 16-row cells, reached only with the + # fused verify). + not_rows = { + dispatch.LONG_SELECT.path, + dispatch.LONG_SELECT_DSL.path, + dispatch.UGPU_HEAD.path, + dispatch.INDEX_SCORE.path, + } + for path in dispatch.FAST_PATHS: + if path.path in raised or path.path in not_rows: + continue + for key in path.keys: + if isinstance(key, int): + assert key <= dispatch.TINY_MAX_TOKENS, (path.path, key) + + +def test_the_batch_4_moe_finalizes_on_the_multicast_kernel() -> None: + """1..8 rows keep the MoE engine; 9..16 rows (past it) take the multicast finalize + all-reduce + norm on the + 16-token bucket, or the generic MoE when B4_MC_FINALIZE is killed; nothing past the decode cap.""" + for rows in range(1, dispatch.TINY_MAX_TOKENS + 1): + assert dispatch.moe_path(rows) == ( + dispatch.MOE_ENGINE.path if not dispatch.KILL_MOE_ENGINE else None + ), rows + for rows in range(dispatch.TINY_MAX_TOKENS + 1, dispatch.DECODE_MAX_TOKENS + 1): + expected = dispatch.MC_FINALIZE.path if dispatch.B4_MC_FINALIZE else None + assert dispatch.moe_path(rows) == expected, rows + assert dispatch.moe_path(dispatch.DECODE_MAX_TOKENS + 1) is None + + +# ------------------------------------------------------------------ W-G1: the MXFP8 GEMV at 33..48 rows (agg wide +# decode) + +_W48 = (dispatch.QKV_SPARSE, dispatch.QKV, dispatch.SHARED_GATE_UP, dispatch.SHARED_DOWN) + + +def test_w48_rows_follow_the_kill_constant() -> None: + """WIDE_MX_GEMV_48 raises the four projections it wins on to their served rows past 32 (at most 48); o_proj and the + dense MLP keep theirs.""" + on = ( + dispatch.MX_GEMV_TP2 + and dispatch.WIDE_MX_GEMV_48 + and dispatch.GEMV_MAX_TOKENS == dispatch.DECODE_MAX_TOKENS + ) + for nk in _W48: + rows = dispatch._GEMV_W48_ROWS.get(nk, frozenset()) + if not on: + assert not rows, nk + continue + assert rows and min(rows) == dispatch.GEMV_MAX_TOKENS + 1, nk + assert max(rows) <= dispatch.WIDE_GEMV_MAX_TOKENS, nk + assert rows == frozenset(range(min(rows), max(rows) + 1)), nk + assert rows <= dispatch.MX_GEMV_ROWS[nk], nk + for nk in (dispatch.O_PROJ, dispatch.DENSE_GATE_UP, dispatch.DENSE_DOWN): + assert nk not in dispatch._GEMV_W48_ROWS and nk not in dispatch.MX_GEMV_QUANT_ROWS, nk + + +@pytest.mark.parametrize("nk", _W48, ids=str) +def test_w48_bf16_activation_takes_quantize_plus_gemv(nk) -> None: + """A bf16 activation on rows the bf16 GEMV does not serve but the MXFP8 GEMV does: linear_path returns the GEMV + (M3MXFP8LinearMethod.apply quantizes with the decode-graph quantize first); rows the bf16 GEMV serves keep it; rows + past the MXFP8 GEMV's reach keep the generic GEMM.""" + if not (dispatch.MX_GEMV_TP2 and dispatch.WIDE_MX_GEMV_48): + return + bf16_hi = dispatch.MX_GEMV_BF16_MAX_TOKENS.get(nk, 0) + for rows in range(1, 65): + path = dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=False) + if rows <= bf16_hi: + assert path == dispatch.MX_GEMV_BF16_PATH, rows + elif rows in dispatch.MX_GEMV_ROWS[nk]: + assert rows in dispatch.MX_GEMV_QUANT_ROWS[nk], rows + assert path == dispatch.MX_GEMV_PATH, rows + else: + assert path == dispatch.GENERIC_GEMM_PATH, rows + if rows in dispatch.MX_GEMV_ROWS[nk]: + assert ( + dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=False) + == dispatch.MX_GEMV_PATH + ), rows + + +def test_w48_rows_are_certified_by_the_gemv_cells() -> None: + """Every (projection, rows) the dispatch can send to the MXFP8 GEMV is one of m3_mx_gemv's CELLS (served_cells to + 48).""" + cells = {(c.spec["nk"], c.spec["m"]) for c in dispatch.m3_mx_gemv.CELLS} + for nk, rows in dispatch.MX_GEMV_ROWS.items(): + for m in rows: + assert (nk, m) in cells, (nk, m) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py new file mode 100644 index 000000000000..a7dba3ce84f3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention under a breakable prefill CUDA graph (``M3Eagle3Draft.attention``): it runs +in the eager draft attention boundary on the step's metadata, read at the call (as the in-tree drafter's attention +boundary reads it), and writes the live rows of the padded output; eagerly it calls the op as before. A metadata that +is not the step's is refused. CPU only: the attention op, the linears and the reduction are stand-ins.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) + +pytestmark = pytest.mark.cpu_only + +Q_SIZE = 8 + + +@pytest.fixture +def draft(monkeypatch): + """(drafter, calls of the attention op stand-in, the step's metadata).""" + calls = [] + step_md = SimpleNamespace(num_tokens=3) + + def op(attn, qkv, md): + calls.append((attn, qkv, md)) + return torch.full((int(md.num_tokens), Q_SIZE), 7.0, dtype=torch.bfloat16) + + monkeypatch.setattr(modeling, "m3_draft_attention", op) + monkeypatch.setattr(modeling, "_record", lambda *a: None) + monkeypatch.setattr( + modeling, "get_model_extra_attrs", lambda: {"attention_metadata": lambda: step_md} + ) + d = object.__new__(modeling.M3Eagle3Draft) + d.__dict__.update(attn="attn", q_size=Q_SIZE, _bcg_key="draft-under-test") + d.linear = lambda x, name: x + monkeypatch.setitem(modeling._BCG_DRAFTS, d._bcg_key, d) + return d, calls, step_md + + +def test_breakable_graph_runs_the_eager_boundary_on_the_steps_metadata(draft, monkeypatch) -> None: + d, calls, step_md = draft + monkeypatch.setattr(modeling, "is_in_breakable_cuda_graph", lambda: True) + x = torch.zeros(5, 12, dtype=torch.bfloat16) # 3 live rows + 2 pad rows of the captured bucket + out = modeling.M3Eagle3Draft.attention(d, x, step_md, False) + assert len(calls) == 1 and calls[0][0] == "attn" and calls[0][2] is step_md + assert out.shape == (5, Q_SIZE) + assert torch.equal(out[:3], torch.full((3, Q_SIZE), 7.0, dtype=torch.bfloat16)) + + +def test_breakable_graph_refuses_a_metadata_that_is_not_the_steps(draft, monkeypatch) -> None: + d, calls, _ = draft + monkeypatch.setattr(modeling, "is_in_breakable_cuda_graph", lambda: True) + with pytest.raises(RuntimeError, match="metadata is not the step's"): + modeling.M3Eagle3Draft.attention( + d, torch.zeros(3, 12), SimpleNamespace(num_tokens=3), False + ) + assert not calls + + +def test_eager_calls_the_op_directly(draft, monkeypatch) -> None: + d, calls, _ = draft + monkeypatch.setattr(modeling, "is_in_breakable_cuda_graph", lambda: False) + md = SimpleNamespace(num_tokens=2) + out = modeling.M3Eagle3Draft.attention(d, torch.zeros(2, 12), md, False) + assert len(calls) == 1 and calls[0][2] is md and out.shape == (2, Q_SIZE) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py new file mode 100644 index 000000000000..7da5d1206a57 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter's full-context decode (M3DrafterFullCtxFmha) is selected through FmhaManager, which caches a library's +support per cell of the generation request count (a ceiling bucket of _FMHA_CACHE_BATCH_SIZE_GRID, e.g. 41..47 -> 47). +The runner supports exact (query rows, requests) recipes, so every recipe's request count must be alone in its cell: +otherwise the first count seen in a cell decides for the others (a 46-request recipe would share 47's cell, and an +eager 41-request pass would reach the runner without a plan). CPU only.""" + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, +) +from tensorrt_llm._torch.attention.backends.fmha import manager + +pytestmark = pytest.mark.cpu_only + + +def _cell(value: int, grid) -> int: + return manager._normalize_fmha_cache_grid_value(value, grid) + + +def test_every_drafter_recipe_is_alone_in_its_cache_cells(): + keys = set(modeling.FULLCTX_DECODE_RECIPES) + batch, qlen = manager._FMHA_CACHE_BATCH_SIZE_GRID, manager._FMHA_CACHE_SEQ_LEN_Q_GRID + for q, n in keys: + same_batch = {m for m in range(1, 4 * batch[-1]) if _cell(m, batch) == _cell(n, batch)} + same_qlen = {r for r in range(1, 4 * qlen[-1]) if _cell(r, qlen) == _cell(q, qlen)} + assert same_batch == {n}, ( + f"({q}, {n}): requests {sorted(same_batch)} share its FmhaManager cell" + ) + assert same_qlen == {q}, ( + f"({q}, {n}): query rows {sorted(same_qlen)} share its FmhaManager cell" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..dfb89711db4d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode robustness without a GPU kernel: the multicast all-reduce epilogue options keep the cluster-barrier bit, +an engine build that fails on one TP rank drops that engine on every rank (two ranks simulated with threads), and +the Eagle3 fc fork's output is recorded on the stream that reads it.""" + +import threading +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = modeling.require_mode() + yield + modeling.set_require(require) + modeling.reset_logged() + + +@pytest.mark.parametrize("finalize", [False, True]) +@pytest.mark.parametrize("bucket", [8, 16, 32, 48]) +def test_ar_opts_keep_cluster_barrier_bit(monkeypatch, bucket, finalize): + # Tables without bit 2 and with buckets missing: every lookup, fallback included, keeps the barrier. + monkeypatch.setattr(mc, "AR_OPTS_BY_M", (0, 1, 4, 5, 8, 9, 12, 13)) + monkeypatch.setattr(mc, "AR_OPTS_BY_BUCKET", {}) + monkeypatch.setattr(mc, "AR_FIN_OPTS_BY_BUCKET", {}) + for m in range(1, min(bucket, 8) + 1): + opts = mc._ar_opts(m, bucket, finalize=finalize) + assert opts & mc._OPT_CLUSTER_BARRIER and 0 <= opts <= 15, (m, bucket, finalize, opts) + + +def test_ar_opts_shipped_tables_unchanged(): + for m in range(1, 9): + assert mc._ar_opts(m, 8) == mc.AR_OPTS_BY_M[m - 1] + for b, o in mc.AR_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b) == o + for b, o in mc.AR_FIN_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b, finalize=True) == o + + +class _FakeComm: + """An MPI-like TP communicator for ``world`` threads: allreduce sums the ranks' ints.""" + + def __init__(self, world): + self.world = world + self.barrier = threading.Barrier(world) + self.lock = threading.Lock() + self.vals = [] + self.calls = 0 + + def allreduce(self, v): + with self.lock: + self.vals.append(v) + self.calls += 1 + self.barrier.wait() + total = sum(self.vals) + self.barrier.wait() + with self.lock: + self.vals.clear() + self.barrier.wait() + return total + + def Get_size(self): # noqa: N802 (MPI name) + return self.world + + +def _mapping(tp_size=2, adp=False): + return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) + + +def _run_ranks(fns): + out, errs = [None] * len(fns), [] + + def run(i): + try: + out[i] = fns[i]() + except BaseException as e: # noqa: BLE001 + errs.append(e) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errs, errs + return out + + +@pytest.fixture +def two_rank_comm(monkeypatch): + from tensorrt_llm._torch.distributed import ops + + comm = _FakeComm(2) + monkeypatch.setattr(mc, "_tp_comm", lambda mapping: comm) + monkeypatch.setattr(ops, "mpi_disabled", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return comm + + +def test_tp_agree_single_rank_and_adp_take_no_collective(monkeypatch): + monkeypatch.setattr(mc, "_tp_comm", mock.Mock(side_effect=AssertionError("no collective"))) + assert mc.tp_agree(_mapping(1), True) is True + assert mc.tp_agree(_mapping(1), False) is False + assert mc.tp_agree(_mapping(4, adp=True), False) is False + + +def test_tp_agree_refuses_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError): + mc.tp_agree(_mapping(2), True) + + +@pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) +def test_tp_agree_built_two_ranks(two_rank_comm, oks): + built = [object() if ok else None for ok in oks] + out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + assert out == (built if all(oks) else [None, None]) + + +def _core(fail: bool) -> SimpleNamespace: + """The engine-building half of the core (``MiniMaxM3Tp2ModelingV2``) on one of two simulated TP ranks, its + engine builds stubbed (and failing on one rank).""" + core = SimpleNamespace(mapping=_mapping(2), moe_engines={}) + cls = modeling.MiniMaxM3Tp2ModelingV2 + + def build(i, num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + core._build_engine = build + for name in ("_prebuild_engines", "_try_build", "engine", "_lookup"): + setattr(core, name, getattr(cls, name).__get__(core)) + return core + + +@pytest.mark.parametrize("fail", [(False, False), (False, True)]) +def test_engine_build_agreed_across_ranks(two_rank_comm, fail): + cores = [_core(f) for f in fail] + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + keys = sorted(cores[0].moe_engines[3]) + assert keys and two_rank_comm.calls == 2 * len(keys) + a, b = cores[0].moe_engines[3], cores[1].moe_engines[3] + assert a.keys() == b.keys() + for key in a: + assert (a[key] is None) == (b[key] is None) == any(fail), key + # Both ranks then take the same path in the forward: the engine, or (after a miss) the generic MoE. + modeling.set_require(False) + for n in sorted(modeling.MOE_ENGINE.keys): + got = [c.engine(3, n, shared=True) for c in cores] + assert (got[0] is None) == (got[1] is None) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + M3Eagle3Worker = modeling.M3Eagle3Worker + main, fork = mock.Mock(name="main"), mock.Mock(name="fork") + hidden = mock.Mock(name="hidden", is_cuda=True) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *a: main) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + worker = SimpleNamespace(_m3_fc_hidden=hidden) + M3Eagle3Worker._m3_fc_join(worker, fork) + main.wait_stream.assert_called_once_with(fork) + if capturing: + hidden.record_stream.assert_not_called() + else: + hidden.record_stream.assert_called_once_with(main) + worker._m3_fc_hidden = None + M3Eagle3Worker._m3_fc_join(worker, fork) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py new file mode 100644 index 000000000000..dcd2ffb2c38b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The wide target's MXFP8 o_proj hand-off, and so its long-context sparse decode, stops at one 128-row scale tile. + +LC2 is certified as the MXFP8 hand-off step, so the hand-off's row cap is LC2's row cap. A wide decode step past it +(33..48 requests x 4 verify rows) takes the Triton producer + merge and the o_proj quantizes its BF16 input itself +(attention/test_modeling_v2_m3_sparse_decode.py: test_plain_route_is_bitwise_the_stock_decode). +""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( + triton_sparse_decode as tsd, +) + + +def test_the_hand_off_cap_is_one_scale_tile(): + """The MXFP8 hand-off (and so the long-context sparse decode) serves up to one 128-row scale tile; the target's + 129..192-row steps take the Triton producer + merge and the o_proj quantizes its BF16 input itself.""" + assert tsd.MergeMxfp8.MAX_ROWS == 128 + assert dispatch.WIDE_MAX_TOKENS > tsd.MergeMxfp8.MAX_ROWS + + +def test_lc2_certifies_only_the_hand_off_step(): + assert dispatch.SPARSE_DECODE_LC.covers("mxfp8") + assert not dispatch.SPARSE_DECODE_LC.covers("plain") + + +@pytest.mark.parametrize("rows", [1, 4, 127, 128]) +def test_the_hand_off_is_the_steps_own_scale_tiles(rows): + """The published scales are the stock MXFP8 quantization's size for the step (the o_proj GEMM checks the swizzled + scales against the step's padded rows): one 128-row tile for every step the hand-off serves.""" + width = 4096 + buffers = tsd.MERGE_MXFP8.buffers(width, torch.device("cuda")) + q8, sf = tsd.MergeMxfp8.handoff(buffers, rows) + _, stock_sf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(rows, width, dtype=torch.bfloat16, device="cuda"), True + ) + assert q8.shape == (rows, width) + assert sf.numel() == stock_sf.numel() + assert sf.data_ptr() == buffers[1].data_ptr() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..b58bcfacd648 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_silu_and_mul entry of the MiniMax-M3 tp4 target: SwiGLU-OAI and plain SwiGLU.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.activation.m3_silu_and_mul import ( # noqa: E501 + m3_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + # Gate and up spread past the clamp (|x| up to ~12) so both clamps engage. + x = (torch.randn(spec["m"], 2 * spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + return (x, *OAI) if spec["oai"] else (x, None, None, None) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few + fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); + measured on GB300: worst element 0.5 unit, relative RMS 0.01.""" + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(*args) + assert out.shape == (cell.spec["m"], cell.spec["inter"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=1.0, rms_ulp=0.1) + + +def test_the_gate_sees_a_missing_clamp() -> None: + """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" + spec = op.CELLS[0].spec + x, *oai = _build(dict(spec, m=64), 3) + out = op(x, *oai) + with pytest.raises(AssertionError): + assert_within_ulp(out, op.reference(x, None, oai[1], oai[2]), element_ulp=1.0, rms_ulp=0.1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_attn_runtime_tp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_attn_runtime_tp4.py new file mode 100644 index 000000000000..7bfd8de384f9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_attn_runtime_tp4.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real MiniMax-M3 attention runtime state at the TP4 per-rank geometry, for the backend-step entries' GPU tests. + +One hybrid NVFP4 KV cache manager (``MiniMaxM3KVCacheManagerV2``) laid out as the target's: a dense layer in the FP8 +half, a sparse MSA layer in the NVFP4 half with its FP8 index-K cache, and (``draft=True``) the one-model Eagle3 +drafter's shared layer, read through the manager's FP8 draft view as the runtime's draft swap reads it. Metadata is +the runtime's own type, filled by its own ``prepare()``; the backends are built the way the target builds them. + +The geometry is one rank's: 16 query heads over one KV head and one index head of dim 128. The mapping is a single +rank holding exactly those heads, so nothing is sharded again. + +History rows are written by the in-tree scatters (``fused_write_layer_caches`` for the FP8 layers, +``fused_write_layer_caches_nvfp4`` for the NVFP4 layer and its index-K), at the slots the manager assigned, which are +not in request order. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + MiniMaxM3MsaSparseAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, + fused_write_layer_caches_nvfp4, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HEADS = 16 +KV_HEADS = 1 +INDEX_HEADS = 1 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +DENSE, SPARSE, DRAFT = 0, 1, 2 +MAX_BATCH = 8 +MAX_SEQ_LEN = 8192 +MAX_NUM_TOKENS = 8192 +#: The Eagle3 drafter checkpoint's RoPE: NeoX over the whole head, theta 5e6. +DRAFT_ROPE_THETA = 5.0e6 + +#: E4M3-exact planted index levels, 1/8 apart in relative terms (0.25 .. 15): 48 distinct block scores. +LEVELS = tuple(2.0**e * (1 + m / 8) for e in range(-2, 4) for m in range(8)) +#: The index value of a block's non-peak rows; its score (8) is below every level's (>= 16). +LOW = 0.125 + + +def mapping() -> Mapping: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + +def sparse_attention_config() -> MiniMaxM3SparseAttentionConfig: + """The target's runtime: MSA, the FP8 indexer, the fused five-way projection (M3's checkpoint selection rule: + no forced initial blocks, the query's own block forced).""" + return MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_num_index_heads=INDEX_HEADS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + ) + + +def make_manager(*, draft: bool) -> MiniMaxM3KVCacheManagerV2: + torch.cuda.init() + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_tokens=4 * MAX_SEQ_LEN, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + event_buffer_max_size=0, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=[KV_HEADS] * (3 if draft else 2), + head_dim=HEAD_DIM, + tokens_per_block=PAGE, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + mapping=mapping(), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") if draft else None + ), + sparse_layer_ids=[SPARSE], + sparse_attention_config=sparse_attention_config(), + ) + + +def make_msa_backend(layer: int) -> MiniMaxM3MsaSparseAttention: + """One layer's MSA backend as ``MiniMaxM3Tp4ModelingV2`` builds it (NVFP4 KV quant config on every layer).""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + backend = create_attention( + "TRTLLM", + layer, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=None, + quant_config=quant, + q_scaling=1.0, + sparse_params=sparse_attention_config().to_sparse_params(layer_idx=layer), + ) + backend.update_quant_config(quant) + assert isinstance(backend, MiniMaxM3MsaSparseAttention), type(backend).__name__ + return backend + + +def make_draft_backend() -> TrtllmAttention: + """The drafter layer's attention as ``M3Eagle3Draft`` builds it: fused NeoX RoPE, an FP8 KV cache (the target + overrides the drafter's KV dtype to FP8), no other quantization.""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + backend = create_attention( + "TRTLLM", + DRAFT, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=HEAD_DIM, theta=DRAFT_ROPE_THETA, max_positions=MAX_SEQ_LEN), + ), + quant_config=quant, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + backend.update_quant_config(quant) + assert type(backend) is TrtllmAttention, type(backend).__name__ + return backend + + +def make_msa_metadata(manager) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = MiniMaxM3MsaSparseAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + max_num_sequences=MAX_BATCH, + kv_cache_manager=manager, + mapping=mapping(), + sparse_metadata_params=sparse_attention_config().to_sparse_metadata_params(), + ) + # The target's MSA preparation choice that changes what prepare() stages (the identity tables of short steps). + metadata.msa_policy = MsaRuntimePolicy(short_sequence_identity=True) + return metadata + + +def make_draft_metadata(manager) -> TrtllmAttentionMetadata: + """Drafter metadata over the shared manager's FP8 draft view, the manager the runtime's draft swap installs.""" + view = manager.get_draft_kv_cache_view() + assert view is not None and view.dtype == DataType.FP8 + return TrtllmAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=mapping(), + ) + + +def allocate(manager, request_ids: Sequence[int], qlens: Sequence[int], cached: Sequence[int]): + manager.add_dummy_requests(list(request_ids), token_nums=[c + n for c, n in zip(cached, qlens)]) + + +def prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + c + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +def flat_slots(manager, layer: int, request_id: int, positions: torch.Tensor) -> torch.Tensor: + """The cache slots (page slot id * PAGE + offset) of a request's positions in ``layer``.""" + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0], + dtype=torch.int64, + device="cuda", + ) + positions = positions.to(device="cuda", dtype=torch.int64) + return (pages[positions // PAGE] * PAGE + positions % PAGE).to(torch.int32) + + +def write_fp8_rows(manager, layer, request_id, positions, k, v) -> None: + """K / V rows [n, 128] (E4M3) into an FP8 layer's cache at ``positions``.""" + view = manager.get_buffers(layer, "HND") + assert view.dtype == torch.float8_e4m3fn + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches(view[:, 0], view[:, 1], None, slots, k, v, None) + + +def write_nvfp4_rows(manager, metadata, layer, request_id, positions, k, v, idx_k, quant) -> None: + """BF16 K / V rows quantized into the NVFP4 layer's cache, and index-K rows into its FP8 index cache.""" + data = manager.get_buffers(layer, "HND") + scales = manager.get_block_scale_buffers(layer, "HND") + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches_nvfp4( + data[:, 0], + data[:, 1], + scales[:, 0], + scales[:, 1], + metadata.msa_idx_k_cache(layer), + slots, + k, + v, + idx_k, + quant, + ) + + +def index_directions(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """Two orthogonal ±1 index directions (the halves of the 128 dims), E4M3-exact.""" + signs = torch.randint(0, 2, (HEAD_DIM,), generator=gen, device="cuda").float() * 2 - 1 + u0 = torch.zeros(HEAD_DIM, device="cuda") + u1 = torch.zeros(HEAD_DIM, device="cuda") + u0[: HEAD_DIM // 2] = signs[: HEAD_DIM // 2] + u1[HEAD_DIM // 2 :] = signs[HEAD_DIM // 2 :] + return u0, u1 + + +def planted_index_keys(length: int, u0, u1, gen: torch.Generator) -> torch.Tensor: + """Index-K rows [length, 128] (bf16, E4M3-exact) whose per-block maximum score is unambiguous. + + Every block has one peak row at a random position, ``a_b u0 + c_b u1`` with ``a`` and ``c`` two independent + permutations of ``LEVELS``; its other rows are ``LOW (u0 + u1)``. A query along ``u0`` then ranks the blocks by + ``a``, one along ``u1`` by ``c``, every score 1/8 or more apart from the next.""" + blocks = (length + PAGE - 1) // PAGE + assert blocks <= len(LEVELS) + levels = torch.tensor(LEVELS, device="cuda") + a = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + c = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + rows = (LOW * (u0 + u1)).expand(length, HEAD_DIM).clone() + for b in range(blocks): + span = min(PAGE, length - b * PAGE) + peak = b * PAGE + int(torch.randint(0, span, (1,), generator=gen, device="cuda")) + rows[peak] = a[b] * u0 + c[b] * u1 + return rows.to(torch.bfloat16) + + +def nvfp4_quant(*values: torch.Tensor) -> torch.Tensor: + """The (Q, K, V) quantization scales ``fused_write_layer_caches_nvfp4`` takes: each role's amax at the top of the + E2M1 x E4M3 range.""" + quant = torch.ones(3, dtype=torch.float32, device="cuda") + for role, x in enumerate(values, start=1): + quant[role] = (448.0 * 6.0) / x.abs().max().float() + return quant diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_nvfp4_cache.py new file mode 100644 index 000000000000..2254f6c0b68c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_nvfp4_cache.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A MiniMax-M3 TP4 sparse-layer decode step over a packed NVFP4 cache, for the attention entries' tests. + +The cache is written by the in-tree scatter (``fused_write_layer_caches_nvfp4``, +which quantizes with ``torch.ops.trtllm.fp4_quantize``), into one +``[pages, 2 (K, V), 1, 128, .]`` pool per role exactly as ``msa_paged_kv`` and +``get_block_scale_buffers`` hand the K and V views over. Pages are shuffled so +a kernel that ignored the block table would read the wrong ones. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 + + +def token_kv_len(seq_len: int, decode_query_len: int, offset: int) -> int: + """The KV extent query token ``offset`` of a request attends, as every decode kernel derives it.""" + return max(seq_len - decode_query_len + offset + 1, 0) + + +def build_step( + seq_lens: Sequence[int], decode_query_len: int, *, identity: bool, seed: int +) -> SimpleNamespace: + """One decode step: E4M3 q [T, 16, 128], the packed NVFP4 K/V views and their scales, a [1, T, 16] + head-major top-k table (the identity table, or a random ascending selection with a -1 tail), the block + table and per-request lengths.""" + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = max(1, max((s + PAGE - 1) // PAGE for s in seq_lens)) + num_pages = batch * max_blocks + num_slots = num_pages * PAGE + block_table = ( + torch.randperm(num_pages, device="cuda", generator=g) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + data_pool = torch.zeros(num_pages, 2, 1, PAGE, HEAD_DIM // 2, dtype=torch.uint8, device="cuda") + scale_pool = torch.zeros( + num_pages, 2, 1, PAGE, HEAD_DIM // 16, dtype=torch.uint8, device="cuda" + ) + quant = torch.ones(3, dtype=torch.float32, device="cuda") + values = [] + for role in (0, 1): + x = torch.randn(num_slots, HEAD_DIM, device="cuda", generator=g, dtype=torch.float32).to( + torch.bfloat16 + ) + quant[role + 1] = (448.0 * 6.0) / x.abs().max().float().item() + values.append(x) + assert fused_write_layer_caches_nvfp4( + data_pool[:, 0], + data_pool[:, 1], + scale_pool[:, 0], + scale_pool[:, 1], + None, + torch.arange(num_slots, device="cuda", dtype=torch.int32), + values[0], + values[1], + None, + quant, + ) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) + topk = torch.full((1, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + nblocks = (token_kv_len(seq_lens[req], decode_query_len, offset) + PAGE - 1) // PAGE + if identity: + assert nblocks <= TOPK, "an identity step has every request within 16 blocks" + chosen = torch.arange(nblocks, device="cuda") + else: + chosen = torch.randperm(nblocks, device="cuda", generator=g)[:TOPK].sort().values + topk[0, tok, : chosen.numel()] = chosen.to(torch.int32) + return SimpleNamespace( + q=q, + k_paged=data_pool[:, 0], + v_paged=data_pool[:, 1], + k_block_scale=scale_pool[:, 0], + v_block_scale=scale_pool[:, 1], + k_global_scale=(1.0 / quant[1:2]).contiguous(), + v_global_scale=(1.0 / quant[2:3]).contiguous(), + topk=topk, + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device="cuda"), + decode_query_len=decode_query_len, + sm_scale=HEAD_DIM**-0.5, + source=values, + quant=quant, + num_slots=num_slots, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py new file mode 100644 index 000000000000..824cb6488e58 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_dense_decode entry: the dense layers' trtllm-gen decode at the KV split this target +sizes per request count. + +Every cell runs against the fp32 reference, and a capped cell also against the same call at the model maximum (the +split the in-tree step hands it): the two differ only by the order the CTAs' partial softmax sums are merged in, held +here to a few BF16 ulps. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 + CAPPED, + HEAD_DIM, + HEADS, + KV_HEADS, + MODEL_MAX, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 + m3_dense_decode as op, +) +from tensorrt_llm._torch.attention.backends.fmha.msa_decode import dense_decode_max_seq_len + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_dense_decode requires a CUDA device" + +PAGE = 32 +SUBPAGES = 2 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the kernel asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor): + self.pool = pool + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, SUBPAGES + + +def _inputs(kv_lens, qlen: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * SUBPAGES, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + # Slots handed out non-contiguously, as the block manager does. + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1: a third of each row's softmax mass sits there, far + # outside the band, so a split that missed the tail of a context longer than its max_seq_len fails. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + first_full = (start + PAGE - 1) // PAGE + last = (n + PAGE - 1) // PAGE + if first_full < last: + v_idx = block_table[b, first_full:last].long() * SUBPAGES + 1 + pool[v_idx] += 1.0 + if start % PAGE: + v0 = int(block_table[b, start // PAGE]) * SUBPAGES + 1 + pool[v0, :, start % PAGE :] += 1.0 + q = torch.randn((batch * qlen, HEADS, HEAD_DIM), device="cuda", generator=g).to(torch.bfloat16) + return ( + q, + _Pool(pool.to(torch.float8_e4m3fn)), + block_table, + torch.tensor(kv_lens, device="cuda", dtype=torch.int32), + ) + + +def _args(q, manager, block_table, seq_lens, qlen, max_seq_len): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + return ( + q, + manager, + 0, + block_table, + seq_lens, + SM_SCALE, + out, + qlen, + max_seq_len, + seq_lens.shape[0], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], spec["qlen"], zlib.crc32(cell.why.encode()) % 2**31 + ) + args = _args(q, manager, block_table, seq_lens, spec["qlen"], spec["max_seq_len"]) + certifying(op) + ref = op.reference(*args) + out = op(*args).clone() + torch.testing.assert_close(out.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + if spec["capped"]: + uncapped = op(*_args(q, manager, block_table, seq_lens, spec["qlen"], MODEL_MAX)).clone() + # The split reorders only the merge of the CTAs' fp32 partial sums before the one BF16 rounding: every element + # is within 4 of its own BF16 spacing (2^(e - 8) for |x| = m * 2^e, m in [0.5, 1)). + ref_bf16 = uncapped.float() + _, exp = torch.frexp(ref_bf16.abs()) + spacing = torch.where( + ref_bf16 == 0, + torch.full_like(ref_bf16, 2.0**-133), + torch.ldexp(torch.ones_like(ref_bf16), exp - 8), + ) + assert bool(((out.float() - ref_bf16).abs() <= 4 * spacing).all()) + assert_within_ulp(out, uncapped, element_ulp=4.0, rms_ulp=1.0) + + +@pytest.mark.skipif( + not dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH, + reason="the TP4 SM 10.7 target has no dense-decode cap until one is measured and certified on Rubin", +) +def test_the_cap_is_the_dispatch_constant() -> None: + """The entry's capped cells run at the split the target hands trtllm-gen at one request, and the in-tree step + resolves the cap from the attention module only at the request counts it names.""" + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {1: CAPPED} + attn = type( + "Attn", + (), + {"dense_decode_max_seq_len_by_batch": dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH}, + )() + assert dense_decode_max_seq_len(attn, MODEL_MAX, 1) == CAPPED + assert dense_decode_max_seq_len(attn, MODEL_MAX, 2) == MODEL_MAX + assert dense_decode_max_seq_len(attn, 4096, 1) == 4096 + assert dense_decode_max_seq_len(object(), MODEL_MAX, 1) == MODEL_MAX + + +def test_guard_refuses_more_requests_than_the_counters() -> None: + q, manager, block_table, seq_lens = _inputs((3000, 3000), 1, 3) + args = list(_args(q, manager, block_table, seq_lens, 1, CAPPED)) + args[9] = 1 + with pytest.raises(AssertionError, match="counters cover"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..98c6a27fc039 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7): the TP4 dense layers' FP8 q/k/v (minimax_m3_sm107tp4_dense_qk_norm_rope_fp8, head-count +generic, at the TP4 per-rank 16 query heads over 1 KV head) carries the bits of the stock path, the BF16 +fused_qk_norm_rope output cast to E4M3, for every norm / RoPE variant: Q and K are rounded to BF16 before E4M3 (K is +stored in the FP8 cache), V is cast.""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers the trtllm ops) + +assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) + +HQ, HK, HV, D, ROT = 16, 1, 1, 128, 64 + + +def _entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_dense_qk_norm_rope_fp8 import ( # noqa: E501 + m3_dense_qk_norm_rope_fp8, + ) + + return m3_dense_qk_norm_rope_fp8 + + +def test_certified_cells() -> None: + """The entry (ops/attention/m3_dense_qk_norm_rope_fp8) at every cell: its output is its reference's bit for bit; + with ``kv_write`` the fused-insert form also returns the plain form's bytes and the compact Q.""" + op = _entry() + for cell in op.CELLS: + spec = cell.spec + m, hq, hk = spec["m"], spec["q"], spec["kv"] + g = torch.Generator(device="cuda").manual_seed(m * 7 + hq) + qkv = (torch.randn(m, (hq + 2 * hk) * D, device="cuda", generator=g) * 2).to(torch.bfloat16) + pos = torch.randint(0, 1 << 20, (m,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + args = (qkv, hq, hk, hk, ROT, 1e-6, qw, kw, 5e6, True, pos, True) + ref = op.reference(*args) + if spec.get("kv_write"): + pages, tpb = 4, 128 + cache = torch.zeros(pages, 2, hk, tpb, D, device="cuda", dtype=torch.float8_e4m3fn) + slots = torch.randperm(pages * tpb, device="cuda", generator=g)[:m].to(torch.int32) + q_compact = torch.empty(m, hq, D, device="cuda", dtype=torch.float8_e4m3fn) + out = op(*args, cache, slots, m, q_compact) + assert torch.equal( + q_compact.view(m, hq * D).view(torch.uint8), out[:, : hq * D].view(torch.uint8) + ), cell.why + else: + out = op(*args) + assert torch.equal(out.view(torch.uint8), ref.view(torch.uint8)), cell.why + + +@pytest.mark.parametrize("tokens", [1, 4, 8, 4096]) +@pytest.mark.parametrize("use_gemma", [False, True], ids=["rmsnorm", "gemma"]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_dense_fp8_qkv_is_the_bf16_kernel_cast(tokens, use_gemma, is_neox) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8"): + pytest.skip("this build does not carry minimax_m3_sm107tp4_dense_qk_norm_rope_fp8") + g = torch.Generator(device="cuda").manual_seed(tokens * 4 + int(use_gemma) * 2 + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + ours = torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8( + qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, use_gemma + ) + stock = qkv.clone() + torch.ops.trtllm.fused_qk_norm_rope( + stock, HQ, HK, HV, D, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, + 1.0, 0.0, 0.0, 1.0, True, use_gemma, False, 0, 0, + ) # fmt: skip + assert torch.equal(ours.view(torch.uint8), stock.to(torch.float8_e4m3fn).view(torch.uint8)) + + +@pytest.mark.parametrize("tokens,live", [(1, 1), (4, 4), (8, 5), (64, 64)]) +@pytest.mark.parametrize("tokens_per_block", [32, 128]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_fused_kv_insert_stores_the_bytes_of_write_kv_slots_and_the_compact_q( + tokens, live, tokens_per_block, is_neox +) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + g = torch.Generator(device="cuda").manual_seed(tokens * 1000 + tokens_per_block + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + pages, layers, layer = 16, 3, 1 + # The layer's HND view of a multi-layer pool, strided as get_buffers returns it. + raw = torch.randint( + 0, + 256, + (pages, layers, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.randperm(pages * tokens_per_block, device="cuda", generator=g)[:tokens].to( + torch.int32 + ) + slots[live:] = -1 # padded rows own no slot + + def view(t): + return t.view(torch.float8_e4m3fn)[:, layer].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, False) + q_compact = torch.empty(tokens, HQ, D, device="cuda", dtype=torch.float8_e4m3fn) + ours = torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, live, q_compact + ) + plain = torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8(*args) + assert torch.equal(ours.view(torch.uint8), plain.view(torch.uint8)) + assert torch.equal( + q_compact.view(tokens, HQ * D).view(torch.uint8), plain[:, : HQ * D].view(torch.uint8) + ) + + ref_raw = raw.clone() + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D) + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D) + ref = view(ref_raw) + write_kv_slots(ref[:, 0], slots, k, live, layout="HND") + write_kv_slots(ref[:, 1], slots, v, live, layout="HND") + assert torch.equal(ours_raw, ref_raw) + + +def test_the_fused_kv_insert_skips_a_negative_slot() -> None: + """A live row whose slot is the -1 sentinel stores nothing (no write before the pool row).""" + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + tokens, tokens_per_block, pages = 4, 32, 8 + g = torch.Generator(device="cuda").manual_seed(7) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + raw = torch.randint( + 0, + 256, + (pages, 3, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.tensor([5, -1, 40, 77], device="cuda", dtype=torch.int32) + + def view(t): + return t.view(torch.float8_e4m3fn)[:, 1].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, True, pos, False) + torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, tokens + ) + plain = torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8(*args) + keep = torch.tensor([0, 2, 3], device="cuda") + ref_raw = raw.clone() + ref = view(ref_raw) + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D)[keep] + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D)[keep] + write_kv_slots(ref[:, 0], slots[keep], k, 3, layout="HND") + write_kv_slots(ref[:, 1], slots[keep], v, 3, layout="HND") + assert torch.equal(ours_raw, ref_raw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..a594c099c2ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_attention entry: the Eagle3 drafter layer's attention against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp4``): the target's hybrid NVFP4 +``MiniMaxM3KVCacheManagerV2`` with its one-model Eagle3 draft layer, the requests allocated in it, drafter +``TrtllmAttentionMetadata`` over the manager's FP8 draft view (the manager the runtime's draft swap installs) filled +by its ``prepare()``, and the drafter's ``TrtllmAttention`` built as ``M3Eagle3Draft`` builds it (fused NeoX RoPE at +the drafter checkpoint's theta, FP8 KV). Earlier rows go into the draft layer's cache through the in-tree scatter at +the slots the manager assigned; the call rotates and appends the step's own. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp4 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. +# The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: +# it takes the rotated Q in E4M3 (rounding the reference's Q to E4M3 takes the RMS gap from ~9 to ~7 ulps) and rounds +# the softmax probabilities before the PV product. Measured on GB300 over these cells (max element / RMS): 13.6 / 8.6 +# (4 x 4 rows), 9.9 / 7.2 (8 x 1), 20.4 / 7.7 (512 context rows). The band leaves about 1.5x; a reference shifted +# by one row is far outside it. +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 + +# Cached tokens per request of each cell; the new rows per request come from the cell. +_CACHED = { + (4, 4): (200, 129, 1000, 64), + (8, 1): (5, 127, 128, 300, 1000, 2047, 64, 900), + (1, None): (0,), +} +#: Rows past the live ones in the fused QKV, as a padded step hands them over. +_PAD = 3 + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + +# Non-paged context FMHA on SM107 does not apply RoPE (this fixture's path; the served drafter runs the paged context +# FMHA, which rotates): the context cell is a strict expected failure until the runtime rotates there too. +_UNROTATED_CONTEXT = pytest.mark.xfail( + strict=True, + reason="non-paged context FMHA on SM107 does not apply RoPE (shared runtime)", +) +_CELLS = [ + pytest.param(c, marks=_UNROTATED_CONTEXT) if c.spec.get("context") else c for c in op.CELLS +] + + +@pytest.mark.parametrize("cell", _CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + context = spec.get("context") + cached = _CACHED[spec["reqs"], spec.get("qlen")] + qlens = (context,) if context else (spec["qlen"],) * spec["reqs"] + assert len(cached) == spec["reqs"] + contexts = 1 if context else 0 + rids = list(range(spec["reqs"])) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = rt.make_draft_metadata(manager) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + assert int(metadata.num_tokens) == total + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + certifying(op) + out = op(backend, qkv, metadata) + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{cell.why}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + # The call's append: each live row's rotated K and its V at the row's position, E4M3 (unit KV scale), to + # within one E4M3 step of the float64 rotation (the kernel rotates in FP32). + pool = manager.get_buffers(rt.DRAFT, "HND") + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + want_k = _rope(rows[:, q_size : q_size + rt.HEAD_DIM].float(), positions) + want_v = rows[:, q_size + rt.HEAD_DIM :].to(torch.float8_e4m3fn) + got_k = pool[slots // rt.PAGE, 0, 0, slots % rt.PAGE].to(torch.float64) + got_v = pool[slots // rt.PAGE, 1, 0, slots % rt.PAGE] + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), "K is appended rotated, in E4M3" + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose rows are shifted by one token is outside it. + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py new file mode 100644 index 000000000000..6e22a6c33f36 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for kernels/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V +append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row +first draft pass and the 1-row later passes of one request. + +Lengths sweep the KV tile edges each query row in turn ends on (the rows of a Q CTA with no key in its trailing +splits); every call is held to the fp32 reference over the cache after the append (E4M3 query, unit KV scales, no +RoPE) within the FP8-KV band, finite, for both split-KV reducers. +""" + +import math +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_tokmap_guard as tokmap_guard, +) + +assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +FP8_TOL = 6e-2 +# Persistent like the FMHA manager's workspace and the spec-decoding metadata of a real step: the in-tree preprocess +# kernels must not see these move between calls. +_PERSISTENT = {} +# Row counts of the one-request token maps the guard wrote; the model installs it in prebuild_kernels. +_TOKMAP_TAKEN = [] + + +@pytest.fixture(scope="module", autouse=True) +def _tokmap_guard_installed(): + # A one-request spec pass places its rows by the workspace token map, which the op builds only for several + # requests; without the guard the call reads whatever the workspace holds. + tokmap_guard.install(torch.device("cuda"), dispatch.DECODE_MAX_TOKENS, _TOKMAP_TAKEN.append) + yield + + +def _persistent(name: str, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _edge_lens(plan, qlen: int): + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +def _library(qlen: int, nreq: int, own_reducer: bool): + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = own_reducer + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = torch.cuda.get_device_properties(0).multi_processor_count + lib._on_taken = lambda rows: None + # The in-tree preprocess pass (its RoPE-free cache append); the one-launch preprocess has its own bitwise test. + lib._fused_rope_mode = None + lib._pdl = False + lib._on_rope_taken = None + return lib + + +def _attn(qlen: int): + return SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=0, theta=10000.0, scale_type=0, scale=1.0, max_positions=1024 + ), + ) + + +def _run(lib, qlen: int, kv_lens, seed: int, spec: bool = True): + """One drafter generation pass over a fresh FP8 cache; returns (output, fp32 reference). ``spec`` passes the + spec-decoding lengths/positions (a multi-token generation step); without them a multi-row pass is the drafter's + plain generation step of input_seq_length rows, as it runs in-model.""" + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g).clamp_(-4, 4) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + for r, n in enumerate( + kv_lens + ): # shift the tail third of V, so a missed trailing split fails the band + start = 2 * n // 3 + for pg in range(start // PAGE, math.ceil(n / PAGE)): + pool[int(perm[r, pg]), 1, :, (start % PAGE if pg == start // PAGE else 0) :] += 1.0 + pool = pool.to(torch.float8_e4m3fn) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + pointers = torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64) + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=dev, generator=g) + seq = torch.tensor(kv_lens, dtype=torch.int32, device=dev) + one = torch.ones(1, dtype=torch.float32, device=dev) + multi = qlen > 1 and spec + out = torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=_attn(qlen), + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=pointers, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=seq, + spec_decoding_generation_lengths=( + _persistent( + f"gen_lengths{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=( + _persistent( + f"position_offsets{qlen}x{nreq}", + lambda: ( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous() + ), + ) + if multi + else None + ), + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=out, + ) + taken = len(_TOKMAP_TAKEN) + lib.run_generation(params) + torch.cuda.synchronize() + if multi and nreq == 1: + assert _TOKMAP_TAKEN[taken:] == [ntok], ( + "the one-request spec pass ran without its token map" + ) + # fp32 reference over the cache after the append: E4M3 query, unit scales, causal over the step's rows. + q = qkv[:, : H * D].float().view(nreq, qlen, H, D).to(torch.float8_e4m3fn).float() + flat = pool.float() + ref = torch.empty(nreq, qlen, H, D, device=dev) + for r, n in enumerate(kv_lens): + blocks = perm[r, : math.ceil(n / PAGE)] + k = flat[blocks, 0].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + v = flat[blocks, 1].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + k = k.repeat_interleave(H // KVH, dim=0) + v = v.repeat_interleave(H // KVH, dim=0) + s = torch.einsum("qhd,hkd->hqk", q[r], k) * D**-0.5 + visible = ( + torch.arange(n, device=dev)[None, :] + < (n - qlen + 1 + torch.arange(qlen, device=dev))[:, None] + ) + s = s.masked_fill(~visible[None], float("-inf")) + ref[r] = torch.einsum("hqk,hkd->qhd", torch.softmax(s, dim=-1), v) + return out.float().view(nreq, qlen, H, D), ref + + +@pytest.mark.parametrize("qlen,nreq", sorted(dispatch.FULLCTX_DECODE_RECIPES)) +def test_the_drafter_pass_at_tile_edges(qlen, nreq) -> None: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if fullctx.has_own_reducer() else (False,) + for own in reducers: + fullctx.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_LEN, own, get_env_enable_pdl() + ) + lens = _edge_lens(fullctx.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + lens = [n for n in lens if n > qlen] + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for i, kv_lens in enumerate(batches): + # The first draft pass runs in-model as a plain multi-row generation step; every other batch also takes the + # spec-decoding form. + spec = qlen > 1 and i % 2 == 1 + for own in reducers: + got, ref = _run(_library(qlen, nreq, own), qlen, kv_lens, sum(kv_lens), spec=spec) + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +def _drafter_model(scale: float): + """A drafter whose one attention module holds a TrtllmAttention backend as ``attn`` (the backend is no module).""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + backend = TrtllmAttention.__new__(TrtllmAttention) + backend.kv_cache_scaling_factor = torch.full((1,), scale, device="cuda") + backend.is_mla_enable, backend.kv_lora_rank, backend.v_head_dim, backend.head_dim = ( + False, + None, + None, + D, + ) + backend._fmha_manager = SimpleNamespace(fmha_libs=[]) + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.zeros(1, device="cuda")) + model.self_attn = torch.nn.Module() + model.self_attn.attn = backend + return model, backend + + +def test_install_declines_a_scaled_kv_cache() -> None: + """At unit KV scales the arithmetic is the stock decode's; a scaled cache keeps the stock path.""" + reasons = [] + model, layer = _drafter_model(0.5) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 0 + assert layer._fmha_manager.fmha_libs == [] and reasons == ["a scaled FP8 KV cache"] + model, layer = _drafter_model(1.0) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 1 + assert isinstance(layer._fmha_manager.fmha_libs[0], drafter.M3DrafterFullCtxFmha) + + +def test_a_cache_bound_past_the_plan_capacity_keeps_the_stock_path() -> None: + """The plans are compiled for ``max_kv_len`` of K/V: a step whose cache length bound is longer falls through.""" + + from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase + + class _Backend: + num_heads, num_kv_heads, head_dim, quant_mode = H, KVH, D, FP8_KV_QUANT_MODE + kv_scale_orig_quant = kv_scale_quant_orig = None + + backend = _Backend() + backend.kv_scale_orig_quant, backend.kv_scale_quant_orig = object(), object() + reasons = [] + lib = _library(1, 1, False) + lib._on_uncovered = reasons.append + lib._attn_ref = weakref.ref(backend) + fwd = SimpleNamespace( + attention_input_type=None, + kv_scale_orig_quant=backend.kv_scale_orig_quant, + kv_scale_quant_orig=backend.kv_scale_quant_orig, + ) + meta = SimpleNamespace(num_contexts=0, num_generations=1, max_seq_len=MAX_LEN + 8) + q = torch.empty(1, H * D, device="cuda") + assert not lib._is_supported(q, None, None, meta, fwd, phase=FmhaPhase.GENERATION) + assert reasons == ["the cache length bound is beyond the plans' K/V capacity"] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py new file mode 100644 index 000000000000..56d0a258d3ba --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -0,0 +1,371 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for minimax_m3_sm107tp4_drafter_rope_fp8_kv, the TP4 drafter layer's one-launch generation +preprocess: bitwise against the in-tree ``thop.trtllm_gen_generation_preprocess`` followed by the torch E4M3 cast +of its BF16 query, on the query bytes and on every byte of the paged cache, with GPT-NeoX RoPE from a cos/sin table +(full and partial rotary width), multi-token (spec-decoding positions) and plain passes, one and two requests, and +query values past the E4M3 range. The library's fused pass is also held bitwise to its preprocess pass end to end. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_drafter_rope_fp8_kv"): + pytest.skip("minimax_m3_sm107tp4_drafter_rope_fp8_kv is not built", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +ROPE_GPT_NEOX = 2 +_PERSISTENT = {} + + +def _persistent(name, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _case(qlen, kv_lens, rotary_dim, multi, seed): + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = (torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g) * 2).to( + torch.float8_e4m3fn + ) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, device=dev, generator=g) * 4 + # A few query / key values near and past the E4M3 range (both casts saturate at 448). + hot = torch.rand(qkv.shape, device=dev, generator=g) < 0.01 + qkv = torch.where( + hot, qkv.sign() * (430 + 80 * torch.rand(qkv.shape, device=dev, generator=g)), qkv + ) + qkv = qkv.to(torch.bfloat16).contiguous() + max_pos = max(kv_lens) + 64 + cos_sin = torch.randn(max_pos, rotary_dim // 2, 2, device=dev, generator=g).float().contiguous() + attn = SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=ROPE_GPT_NEOX, + rotary_inv_freq=None, + rotary_cos_sin=cos_sin, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=rotary_dim, theta=10000.0, scale_type=0, scale=1.0, max_positions=max_pos + ), + ) + one = torch.ones(1, dtype=torch.float32, device=dev) + spec_offsets = None + if multi: + spec_offsets = _persistent( + f"pos{qlen}x{nreq}", + lambda: torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous(), + ) + # Spec positions are not the plain row index: shift them per request. + spec_offsets.copy_( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1) + + torch.arange(nreq, dtype=torch.int32, device=dev)[:, None] + ) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=attn, + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=torch.tensor(kv_lens, dtype=torch.int32, device=dev), + spec_decoding_generation_lengths=( + _persistent( + f"len{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=spec_offsets, + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev), + ) + return params, pool + + +def _preprocess(params, mpc): + """run_generation's in-tree preprocess call and the E4M3 cast of its query.""" + attn, meta, fwd, rope = params.attn, params.meta, params.fwd, params.attn.rope_params + q_processed = thop.trtllm_gen_generation_preprocess( + params.qkv_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope.dim, + rope.theta, + int(rope.scale_type), + rope.scale, + rope.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + False, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + mpc, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + )[0] + rows = params.num_tokens * attn.num_heads * attn.head_dim + return ( + q_processed.flatten()[:rows] + .to(torch.float8_e4m3fn) + .view(params.num_tokens, attn.num_heads, attn.head_dim) + ) + + +def _op(params, pool, mode): + attn = params.attn + nreq, qlen = params.batch_size, params.input_seq_length + kv_pages = pool.view(-1, KVH, PAGE, D) + tables = params.meta.kv_cache_block_offsets[0, :nreq] + q = torch.empty( + params.num_tokens, + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=pool.device, + ) + multi = params.spec_decoding_generation_lengths is not None and qlen > 1 + torch.ops.trtllm.minimax_m3_sm107tp4_drafter_rope_fp8_kv( + params.qkv_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + attn.rope_params.dim, + PAGE, + mode, + True, + ) + return q + + +# One-request multi-token passes keep the in-tree preprocess (it reads an unbuilt token map), so the spec cases hold +# two requests. +CASES = [ + (qlen, kv_lens, rotary, multi) + for qlen, multi in ((1, False), (4, True), (4, False)) + for kv_lens in (((4, 130) if multi else (4,)) if qlen > 1 else (130,), (8191, 131073)) + for rotary in (128, 64) +] + + +def _bytes(t): + return t.contiguous().view(torch.uint8) + + +def _matches(case, seed): + """Per rope mode: (query bytes equal, cache bytes equal) against the in-tree preprocess + cast.""" + qlen, kv_lens, rotary, multi = case + kv_lens = [max(n, qlen) for n in kv_lens] + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(qlen, kv_lens, rotary, multi, seed) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + pool_ref = pool.clone() + out = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + out[mode] = ( + bool(torch.equal(_bytes(q), _bytes(q_ref))), + bool(torch.equal(_bytes(pool), _bytes(pool_ref))), + ) + return out + + +@pytest.mark.parametrize( + "case", + CASES, + ids=lambda c: f"q{c[0]}-{'x'.join(map(str, c[1]))}-r{c[2]}-{'spec' if c[3] else 'plain'}", +) +def test_the_pinned_rope_mode_is_bitwise(case) -> None: + result = _matches(case, seed=7) + print(f"[rope modes] {case}: {result}", flush=True) + assert result[dispatch.DRAFTER_ROPE_MODE] == (True, True), result + + +def test_the_fused_pass_matches_the_preprocess_pass() -> None: + mpc = torch.cuda.get_device_properties(0).multi_processor_count + for qlen, nreq in sorted(dispatch.FULLCTX_DECODE_RECIPES): + for multi in (False, True) if qlen > 1 else (False,): + kv_lens = [70001 + 517 * r for r in range(nreq)] + outs = [] + for fused in (False, True): + params, pool = _case(qlen, kv_lens, 128, multi, seed=11) + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = True + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = mpc + lib._fused_rope_mode = dispatch.DRAFTER_ROPE_MODE if fused else None + lib._pdl = True + taken = [] + lib._on_rope_taken = lambda rows, what: taken.append(what) + lib._on_taken = lambda rows: None + fullctx.prebuild(torch.device("cuda"), lib._recipes, MAX_LEN, True, True) + lib.run_generation(params) + torch.cuda.synchronize() + assert bool(taken) == (fused and not (multi and nreq == 1)) + outs.append((params.output.clone(), pool.clone())) + assert torch.equal(_bytes(outs[0][1]), _bytes(outs[1][1])), (qlen, nreq, multi, "cache") + assert torch.equal(_bytes(outs[0][0]), _bytes(outs[1][0])), ( + qlen, + nreq, + multi, + "output", + ) + + +def _round_bf16(x): + u = x.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000 + return u.to(torch.int64).to(torch.int32).view(torch.float32) + + +def _discriminating(n, seed): + """(cos, sin, x, pair) quadruples whose rotated first-half output's E4M3 byte differs between the rope modes + (fp32 CPU emulation; the GPU comparison is the verdict).""" + g = torch.Generator().manual_seed(seed) + keep = [] + while sum(k.shape[0] for k in keep) < n: + m = 1 << 20 + c, s = torch.randn(m, generator=g), torch.randn(m, generator=g) + x = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + xp = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + sn = -s + r0 = (c * x) + (sn * xp) + r1 = ((c.double() * x.double()) + (sn * xp).double()).float() + r2 = ((sn.double() * xp.double()) + (c * x).double()).float() + b = [ + _round_bf16(r.contiguous()).to(torch.float8_e4m3fn).view(torch.uint8) + for r in (r0, r1, r2) + ] + diff = (b[0] != b[1]) | (b[0] != b[2]) | (b[1] != b[2]) + diff &= (r0.abs() < 400) & (r1.abs() < 400) + keep.append(torch.stack([c, s, x, xp], 1)[diff]) + return torch.cat(keep)[:n] + + +def test_crafted_inputs_pin_the_rope_mode() -> None: + """Rotations whose E4M3 byte depends on how the two products are summed: the in-tree preprocess and exactly the + pinned mode agree on every one.""" + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(1, [130], 128, False, seed=3) + half = 64 + quads = _discriminating(H * half, seed=5).cuda() + pos = 129 # the one row's rope position: cache length 130, one query row + table = params.attn.rotary_cos_sin.view(-1, half, 2) + # One (cos, sin) per dim: every head shares the position's table row, so the heads take the same dims' quads with + # their own (x, pair). + table[pos, :, 0] = quads[:half, 0] + table[pos, :, 1] = quads[:half, 1] + qkv = params.qkv_input.view(-1, H + 2 * KVH, D) + for h in range(H): + qkv[0, h, :half] = quads[:half, 2].to(torch.bfloat16) + qkv[0, h, half:] = quads[:half, 3].to(torch.bfloat16) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + agree = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + agree[mode] = int((_bytes(q)[0, :, :half] == _bytes(q_ref)[0, :, :half]).sum()) + print( + f"[rope modes crafted] agreeing first-half query bytes of {H * half}: {agree}", flush=True + ) + assert agree[dispatch.DRAFTER_ROPE_MODE] == H * half, agree diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..529a556ea4c3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp8_indexer_qk_norm_rope entry of the MiniMax-M3 tp4 target: index-Q out and index-K in the +cache against the in-place fused norm + RoPE followed by the E4M3 cast, byte for byte up to the 0.1% budget the +in-tree test of the kernel allows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 + m3_fp8_indexer_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" + +DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 + + +def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: + assert a.shape == b.shape and a.dtype == b.dtype == torch.float8_e4m3fn + differ = int((a.view(torch.uint8) != b.view(torch.uint8)).sum()) + assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + m, heads = cell.spec["m"], cell.spec["heads"] + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + qk = torch.randn(m, (heads + 1) * DIM, generator=g, device="cuda").to(torch.bfloat16) + qw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (m,), generator=g, device="cuda", dtype=torch.int32) + pages = (m + PAGE - 1) // PAGE + 1 + cache = torch.zeros(pages, 1, PAGE, DIM, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(pages * PAGE, generator=g, device="cuda")[:m].to(torch.int32) + before = qk.clone() + args = (qk, cache, slots, heads, DIM, ROTARY, EPS, qw, kw, THETA, pos) + with validating(op): + q = op(*args) + assert torch.equal(qk, before), "the BF16 input is left untouched" + q_ref, k_ref = op.reference(*args) + _bytes_close(q.reshape(m, heads, DIM), q_ref) + flat = cache.view(-1, DIM) + _bytes_close(flat[slots.long()], k_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_saturate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_saturate.py new file mode 100644 index 000000000000..a60894191c22 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_saturate.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""paged_cache.fp8_e4m3_saturate (FP8_SATURATE): the E4M3 cache writes and the transient FP8 query casts saturate to ++-448 like the CUDA writers (__NV_SATFINITE); every finite in-range value keeps the plain cast's bits, values already +in E4M3 pass through, and the kill switch is exactly the plain cast. Where torch's cast already saturates (this image) +the conversion is the cast alone, one kernel; otherwise it clamps first, the same bytes.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import paged_cache + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fp8_e4m3_saturate_matches_satfinite(dtype, monkeypatch) -> None: + g = torch.Generator(device="cuda").manual_seed(0) + inside = (torch.randn(4096, generator=g, device="cuda") * 100).clamp(-448, 448).to(dtype) + plain = inside.to(torch.float8_e4m3fn) + assert torch.equal( + paged_cache.fp8_e4m3_saturate(inside).view(torch.uint8), plain.view(torch.uint8) + ) + edge = torch.tensor( + [449.0, 463.0, 464.0, 1e4, float("inf"), -500.0, -float("inf")], dtype=dtype, device="cuda" + ) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0], device="cuda") + assert torch.equal(paged_cache.fp8_e4m3_saturate(edge).float(), want) + nan = paged_cache.fp8_e4m3_saturate(torch.tensor([float("nan")], dtype=dtype, device="cuda")) + assert torch.isnan(nan.float()).all() + already = edge.clamp(-448, 448).to(torch.float8_e4m3fn) + assert paged_cache.fp8_e4m3_saturate(already) is already + monkeypatch.setattr(paged_cache, "FP8_SATURATE", False) + killed = paged_cache.fp8_e4m3_saturate(edge) + assert torch.equal(killed.view(torch.uint8), edge.to(torch.float8_e4m3fn).view(torch.uint8)) + + +def test_cache_write_saturates() -> None: + """The cache-dtype conversion every write_kv_slots branch applies (compared as bytes: torch has no float8 + index_copy_ kernel), and an E4M3 source passing through untouched.""" + values = torch.tensor([[1.0, 500.0, -1e5, float("inf")]], dtype=torch.bfloat16, device="cuda") + out = paged_cache._to_cache_dtype(values, torch.float8_e4m3fn) + assert out.dtype == torch.float8_e4m3fn + want = torch.tensor([[1.0, 448.0, -448.0, 448.0]], device="cuda").to(torch.float8_e4m3fn) + assert torch.equal(out.view(torch.uint8), want.view(torch.uint8)) + assert paged_cache._to_cache_dtype(out, torch.float8_e4m3fn) is out + bf = paged_cache._to_cache_dtype(values, torch.bfloat16) + assert bf.dtype == torch.bfloat16 and torch.equal(bf, values) + + +_EDGES = [ + 448.0, + 449.0, + 463.0, + 464.0, + 1e4, + float("inf"), + -448.0, + -449.0, + -464.0, + -1e4, + -float("inf"), + float("nan"), + 0.0, +] + + +def _cuda_kernels(fn) -> int: + from torch.profiler import ProfilerActivity, profile + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + return sum(1 for e in prof.events() if e.device_type.name == "CUDA") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("cast_saturates", [True, False], ids=["cast_only", "clamp_fallback"]) +def test_fp8_e4m3_saturate_bytes_equal_clamp_then_cast(dtype, cast_saturates, monkeypatch) -> None: + """Both paths give the bytes of clamp(+-448) followed by the cast, edge values included.""" + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", cast_saturates) + x = torch.tensor(_EDGES, dtype=dtype, device="cuda") + want = x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + assert torch.equal(paged_cache.fp8_e4m3_saturate(x).view(torch.uint8), want.view(torch.uint8)) + + +def test_fp8_e4m3_saturate_takes_the_cast_only_path_here(monkeypatch) -> None: + """This image's torch cast saturates, on the CPU (the import-time probe) and on CUDA, so the conversion is one + kernel; with the probe failing it falls back to clamp + cast, two kernels.""" + assert paged_cache._CAST_SATURATES + x = torch.tensor(_EDGES, dtype=torch.bfloat16, device="cuda") + assert torch.equal( + x.to(torch.float8_e4m3fn).view(torch.uint8), + x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).view(torch.uint8), + ) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 1 + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", False) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 2 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py new file mode 100644 index 000000000000..6f2ace58365a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for the attention/m3_fullctx_decode entry: the dense layers' full-context FP8 decode on PrimTS +split-KV decode at the recipe the target pins per (query rows per request, requests). + +Every cell runs against the fp32 reference (q rounded to E4M3 first, as the kernel runs it), eager and replayed from a +CUDA graph (the same bits), and twice (the same bits). The K/V pages are read through the strided views of a flat +sub-page pool, as the KV cache manager hands them out. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_fullctx_decode import ( # noqa: E501 + HEAD_DIM, + HEADS, + KV_HEADS, + MAX_KV_LEN, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_fullctx_decode import ( # noqa: E501 + m3_fullctx_decode as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) + +PAGE = 128 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the entry asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor, subpages: int): + self.pool = pool + self.subpages = subpages + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, self.subpages + + +def _inputs(kv_lens, qlen: int, seed: int, subpages: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * subpages, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1, so a split that missed the tail fails the band. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + for page in range(start // PAGE, (n + PAGE - 1) // PAGE): + v = int(block_table[b, page]) * subpages + 1 + pool[v, :, (start % PAGE if page == start // PAGE else 0) :] += 1.0 + # q as the layer hands it over on decode steps: the FP8 query columns of the fused norm/RoPE kernel's + # [q | k | v] output (row stride wider than q), so no cast makes it compact. + qkv = torch.randn( + (batch * qlen, (HEADS + 2 * KV_HEADS) * HEAD_DIM), device="cuda", generator=g + ).to(torch.float8_e4m3fn) + q = qkv[:, : HEADS * HEAD_DIM].view(batch * qlen, HEADS, HEAD_DIM) + seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32) + return q, _Pool(pool.to(torch.float8_e4m3fn), subpages), block_table, seq_lens + + +@pytest.mark.parametrize("subpages", [2, 3], ids=["kv", "kv+index"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, subpages, certifying) -> None: # noqa: F811 + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], qlen, zlib.crc32(cell.why.encode()) % 2**31, subpages + ) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + certifying(op) + ref = op.reference(*args) + eager = op(*args).clone() + torch.testing.assert_close(eager.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + again = op(*args).clone() + assert torch.equal(again, eager), "not deterministic" + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args) + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager), "the CUDA-graph replay differs from the eager call" + + +# Plans compiled at other K/V stage depths stay alive for the whole session: a compiled decode is cached per config +# and is only ever run with the buffers of a live plan. +_OTHER_DEPTH_PLANS = [] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_kv_stage_depth_keeps_the_bits(cell) -> None: + """kv_stages only sets how many K/V tiles the load pipeline keeps in flight: the recipe's plan and plans compiled + at 12 and at 4 stages give the same bits.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + if not kernel.has_own_reducer(): + pytest.skip("this build does not carry minimax_m3_sm107tp4_fullctx_reduce") + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs(spec["kv_lens"], qlen, 13, 2) + k_pages, v_pages = kernel.kv_page_views(manager.pool, manager.subpages) + kernel.prebuild(q.device, {(qlen, nreq): recipe}, MAX_KV_LEN, True, get_env_enable_pdl()) + plans = [kernel.plan_for(q.device, qlen, nreq, True)] + for stages in (12, 4): + config = {**recipe["config"], "kv_stages": stages} + plans.append( + kernel.FullCtxPlan( + q.device, + qlen, + nreq, + MAX_KV_LEN, + {**recipe, "config": config}, + True, + get_env_enable_pdl(), + ) + ) + _OTHER_DEPTH_PLANS.append(plans[-1]) + outs = [] + for plan in plans: + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + plan.run(q, k_pages, v_pages, block_table, seq_lens, SM_SCALE, out) + outs.append(out) + assert torch.equal(outs[0], outs[1]) and torch.equal(outs[0], outs[2]) + + +def test_the_cells_cover_the_dispatch_rows() -> None: + certified = {(c.spec["qlen"], len(c.spec["kv_lens"])) for c in op.CELLS} + assert set(dispatch.FULLCTX_DECODE.keys) <= certified + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: + """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once + in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + + if not kernel.has_own_reducer(): + pytest.skip("this build does not carry minimax_m3_sm107tp4_fullctx_reduce") + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs(spec["kv_lens"], qlen, 11, 2) + outs = [] + for own in (True, False): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + op(q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe, own_reducer=own) + outs.append(out.float()) + ours, primts = outs + _, exp = torch.frexp(primts.abs()) + spacing = torch.where( + primts == 0, + torch.full_like(primts, 2.0**-133), + torch.ldexp(torch.ones_like(primts), exp - 8), + ) + assert bool(((ours - primts).abs() <= spacing).all()) + + +def _tile_edge_kv_lens(plan, qlen: int): + """KV lengths around every kind of boundary the split ranges have, so each query row of the Q CTA in turn ends + its keys on a KV tile edge (k * tile + d, d in -1 .. qlen - 1): the first tiles, the reducer's 8-group stride, the + lengths where a split's range grows (splits * insts tiles) and a split edge past it.""" + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +@pytest.mark.parametrize("qlen,nreq", sorted(dispatch.FULLCTX_DECODE_RECIPES)) +def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: + """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys + (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode + step first hit (NaN rows from exp2(-inf - -inf)).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, False, get_env_enable_pdl() + ) + lens = _tile_edge_kv_lens(kernel.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, sum(kv_lens), 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +@pytest.mark.parametrize("qlen,nreq", sorted(dispatch.FULLCTX_DECODE_RECIPES)) +def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: + """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just + past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + lengths = [(1 << 20) + 1, (1 << 20) + 8] + ([1_500_001] if nreq == 1 else []) + batches = [[n] * nreq if nreq == 1 else [n, (1 << 20) - 3] for n in lengths] + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, 13, 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, own, get_env_enable_pdl() + ) + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +def _swizzled_scale_offsets(rows: int, width: int) -> torch.Tensor: + """Offsets of rows [0, rows)'s UE8M0 scales in the 128x4-swizzled layout of mxfp8_quantize(x, True).""" + k_sf = width // 32 + k_tiles = (k_sf + 3) // 4 + r = torch.arange(rows, device="cuda").view(-1, 1) + k = torch.arange(k_sf, device="cuda").view(1, -1) + offsets = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + return offsets.reshape(-1) + + +@pytest.mark.parametrize("splits,rows,seq_len_q", [(108, 8, 1), (64, 8, 1), (96, 32, 4)]) +def test_the_merge_mxfp8_copy_is_stock_mxfp8_quantize(splits, rows, seq_len_q) -> None: + """The TP4 merge's MXFP8 copy (the dense o_proj hand-off) equals mxfp8_quantize(out, True) bit for bit: E4M3 rows + and the real rows' swizzled UE8M0 scales, with an all-zero head row (amax 0) and saturating outliers; the BF16 + output is the merge without the copy.""" + ops = torch.ops.trtllm + if not hasattr(ops, "minimax_m3_sm107tp4_fullctx_reduce") or "q8" not in str( + ops.minimax_m3_sm107tp4_fullctx_reduce.default._schema + ): + pytest.skip("needs the TP4 merge op with the MXFP8 copy") + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(splits * 7 + rows) + nreq, kvh = 2, 1 + tokens = nreq * seq_len_q + width = kvh * rows // seq_len_q * 128 + for dtype in (torch.float16, torch.bfloat16): + po = torch.randn(nreq, kvh, splits, rows, 128, device=dev, generator=g) * 3 + po[0, 0, :, 1] = 0.0 + po[1, 0, :, 0, 5] = 6.0e4 + po = po.to(dtype).contiguous() + ps = (torch.randn(nreq, kvh, splits, rows, device=dev, generator=g) * 4).contiguous() + seq = torch.tensor([165521, 120011], dtype=torch.int32, device=dev) + args = (po, ps, seq) + tail = (seq_len_q, seq_len_q, 128, 1, seq_len_q > 1, False) + plain = torch.zeros(tokens, width, dtype=torch.bfloat16, device=dev) + ops.minimax_m3_sm107tp4_fullctx_reduce(*args, plain, *tail) + out = torch.zeros(tokens, width, dtype=torch.bfloat16, device=dev) + q8 = torch.empty(128, width, dtype=torch.float8_e4m3fn, device=dev) + sf = torch.zeros(128 * (width // 32), dtype=torch.uint8, device=dev) + ops.minimax_m3_sm107tp4_fullctx_reduce(*args, out, *tail, q8, sf) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), plain.view(torch.int16)) + stock_q, stock_sf = ops.mxfp8_quantize(out, True) + assert torch.equal( + q8[:tokens].view(torch.uint8), stock_q.view(torch.uint8).reshape(tokens, -1) + ), dtype + at = _swizzled_scale_offsets(tokens, width) + assert torch.equal(sf[at], stock_sf.view(torch.uint8).reshape(-1)[at]), dtype diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..1d4353c466cf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score entry: the MSA index scorer with its pre-wait schedule. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score``) bit for bit, which the pre-wait prefetch and the +early tile issue must not change. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score import ( # noqa: E501 + PAGE, + SPLIT_K, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 1 # index heads per rank at TP4 +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + args = (idx_q, index_k, block_table, seq, score, qlen, SPLIT_K, spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen, SPLIT_K + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +# The 2-request decode step also runs the early-issue kernel on this target (dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS). +# The op file is the sm_103 copy (test_modeling_v2_m3_tp4_drift), so these cells live here. +TWO_REQUEST_EARLY = ( + ((20000, 3000), "two requests verifying 1 + 3 rows each over 20000 and 3000 tokens"), + ( + (70000, 129), + "two requests: one past 512 blocks (its CTAs loop), one whose early tiles are drained", + ), +) + + +def test_the_two_request_step_takes_the_early_kernel() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, + ) + + assert dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS == 2 + + +@pytest.mark.parametrize( + "seq_lens,why", TWO_REQUEST_EARLY, ids=[w[:44] for _, w in TWO_REQUEST_EARLY] +) +def test_two_request_early_issue_is_the_in_tree_scorer(seq_lens, why) -> None: + idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, 4, zlib.crc32(why.encode())) + score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) + args = (idx_q, index_k, block_table, seq, score, 4, SPLIT_K, True) + ref = op.reference(*args) + op(*args) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4, SPLIT_K + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, SPLIT_K, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4, SPLIT_K + ) + assert torch.equal(score, stock) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, SPLIT_K, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py new file mode 100644 index 000000000000..be30baa9bd5b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score_t2sc entry (T2SC, dispatch.INDEX_SCORE_T2SC): the TP2 / TP8 MSA index +scorer with its pre-wait schedule, as the TP4 target's copy at its 1 index head per rank. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score`` at its own 256-CTA split) bit for bit, which the +pre-wait prefetch, the early tile issue and the cell's split must not change. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score_t2sc import ( # noqa: E501 + PAGE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score_t2sc import ( # noqa: E501 + m3_index_score_t2sc as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 1 # index heads per rank at TP4 (4 index heads over 4 KV groups) +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + assert spec["heads"] == HEADS + args = (idx_q, index_k, block_table, seq, score, qlen, spec["split_k"], spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, 1536, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4 + ) + assert torch.equal(score, stock) + + +@pytest.mark.parametrize("qlen", [1, 2, 3, 4]) +def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> None: + """Lengths around the edge between pre-staged (all-old) and post-wait first blocks, 1 and 2 requests.""" + for k in (1, 2, 1294): + for length in (k * PAGE - qlen - 1, k * PAGE - qlen, k * PAGE - qlen + 1, k * PAGE + 1): + for seq_lens, split in ( + ((length,), 1536), + ((length,), 1272), + ((length, max(qlen, length // 3)), 768), + ): + idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) + score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) + op(idx_q, index_k, block_table, seq, score, 4, split, True) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + assert torch.equal(score, stock), (qlen, seq_lens) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, 1536, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..70cbd7b25768 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_msa_attention entry: one layer's MSA backend step against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp4``): a hybrid NVFP4 ``MiniMaxM3KVCacheManagerV2`` +with the requests allocated in it, ``MiniMaxM3MsaSparseAttentionMetadata`` filled by its ``prepare()``, and the layer's +``MiniMaxM3MsaSparseAttention`` built as the target builds it. Earlier rows (and, on a sparse layer, the step's own +rows, which the target's producer writes before the call) go into the cache through the in-tree scatters, at the +slots the manager assigned. The sparse layers' index-K rows are planted so every block's score is 1/8 or more apart +from every other's: the top-k selection the reference recomputes is then unambiguous. + +The cells are the target's calls (``_attention_core`` in modeling.py): a dense layer with live bf16 K / V in a +prefill chunk and live FP8 K / V in a decode step (the dense layers sit in the hybrid cache's FP8 half), a sparse +layer whose producer already wrote the caches (``k = v = None``) in a prefill chunk beyond 16 blocks and in a mixed +step, both with the NVFP4 scales the target passes. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_msa_attention import ( # noqa: E501 + m3_msa_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp4 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" + +# Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, +# per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 +# FMHA does), so what is left is each FMHA's own arithmetic, measured on GB300 over these cells (max element / RMS): +# * fmha_sm100 over the FP8 pool (dense context rows): 10.1 / 6.1. It runs E4M3 x E4M3 and its softmax +# probabilities are rounded before the PV product. +# * trtllm-gen over the FP8 pool (dense generation rows): 7.0 / 6.1. It rounds the probabilities to E4M3 (scaled by +# 448 against the running maximum): emulating that brings a single-tile request to 0.6 RMS. +# * the NVFP4 CSR prefill (sparse context rows): 34.4 / 11.8 at 4096 rows, 26.8 / 9.1 at 512. It stages the +# reconstructed K / V in E4M3 (the in-tree A/B test, test_minimax_m3_sparse_attn_decode, notes it and holds it +# to 3e-2). +# * the Triton NVFP4 sparse decode (sparse generation rows): 1.0 / 0.25 -- the band of its own entry +# (m3_sparse_decode): exact operands, FP16 probabilities, BF16 rounding. +# Each band leaves about 1.5x the measured maximum; a shifted or wrongly selected reference is far outside them. +_FP8_CONTEXT = dict(element_ulp=16.0, rms_ulp=9.0) +_FP8_DECODE = dict(element_ulp=12.0, rms_ulp=9.0) +_NVFP4_CONTEXT = dict(element_ulp=48.0, rms_ulp=16.0) +_NVFP4_DECODE = dict(element_ulp=2.0, rms_ulp=0.5) +_BANDS = { + ("dense", "context"): (_FP8_CONTEXT, None), + ("dense", "generation"): (None, _FP8_DECODE), + ("sparse", "context"): (_NVFP4_CONTEXT, None), + ("sparse", "mixed"): (_NVFP4_CONTEXT, _NVFP4_DECODE), +} + + +def _assert_in_bands(output, ref, bands, context_rows: int) -> None: + """Gate the context rows and the generation rows each on their own kernel's band.""" + context_band, decode_band = bands + if context_rows: + assert_within_ulp(output[:context_rows], ref[:context_rows], **context_band) + if context_rows < output.shape[0]: + assert_within_ulp(output[context_rows:], ref[context_rows:], **decode_band) + + +# (request rows, cached tokens per request, new tokens per request, context rows) per cell. +_SHAPES = { + ("dense", "context"): ((384,), (1024,), 1), + ("dense", "generation"): ((127, 128, 1000, 2500), (1, 1, 1, 1), 0), + ("sparse", "context"): ((0,), (4096,), 1), + ("sparse", "mixed"): ((0, 2999, 4500), (512, 1, 1), 1), +} + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """The two quantities ``assert_within_ulp`` gates, in BF16 ulps.""" + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _shift(x: torch.Tensor) -> torch.Tensor: + """``x``'s rows moved down by one (byte-wise, so E4M3 rows move too).""" + return x.view(torch.uint8).roll(1, 0).view(x.dtype) if x.element_size() == 1 else x.roll(1, 0) + + +def _shape(cell): + spec = cell.spec + cached, qlens, contexts = _SHAPES[spec["layer"], spec["phase"]] + if spec["phase"] == "context": + assert sum(qlens) == spec["tokens"] + elif spec["phase"] == "generation": + assert len(qlens) == spec["reqs"] + else: + assert sum(qlens[:contexts]) == spec["tokens"] and len(qlens) - contexts == spec["reqs"] + return cached, qlens, contexts + + +def _run_dense(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + rt.write_fp8_rows(manager, rt.DENSE, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + # A prefill chunk arrives in BF16 (the hybrid cache's dense layers take FP8 Q / K / V only on decode steps). + dtype = torch.bfloat16 if contexts else torch.float8_e4m3fn + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + k = torch.randn(total, rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + v = torch.randn(total, rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + # The NVFP4 target passes its layer scales on every layer; an FP8 layer's write and FMHA take none. + inv = torch.tensor([1.0, 0.5, 2.0], device="cuda") + fwd = 1.0 / inv + return (q, k, v, None, inv, fwd), dict(rids=rids, cached=cached, qlens=qlens, contexts=contexts) + + +def _run_sparse(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + u0, u1 = rt.index_directions(gen) + lengths = [c + n for c, n in zip(cached, qlens)] + keys = [torch.randn(n, rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths] + values = [torch.randn(n, rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths] + quant = rt.nvfp4_quant(torch.cat(keys), torch.cat(values)) + for rid, n, kk, vv in zip(rids, lengths, keys, values): + # Every row, the step's own included: the producer writes a sparse layer's caches before the call. + idx_k = rt.planted_index_keys(n, u0, u1, gen) + rt.write_nvfp4_rows( + manager, + metadata, + rt.SPARSE, + rid, + torch.arange(n), + kk.to(torch.bfloat16), + vv.to(torch.bfloat16), + idx_k, + quant, + ) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + # Query tokens alternate between the two planted directions, so neighbouring tokens select different blocks. + idx_q = torch.stack([u0, u1])[torch.arange(total, device="cuda") % 2].to(torch.float8_e4m3fn) + return (q, None, None, idx_q, quant, 1.0 / quant), dict( + rids=rids, cached=cached, qlens=qlens, contexts=contexts + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + sparse = spec["layer"] == "sparse" + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=False) + try: + backend = rt.make_msa_backend(rt.SPARSE if sparse else rt.DENSE) + metadata = rt.make_msa_metadata(manager) + run = _run_sparse if sparse else _run_dense + (q, k, v, idx_q, inv, fwd), step = run(cell, gen, manager, backend, metadata) + output = torch.full( + (q.shape[0], rt.HEADS * rt.HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + certifying(op) + result = op(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert result is output + ref = op.reference(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert torch.isfinite(output.float()).all(), "every live row is written" + bands = _BANDS[spec["layer"], spec["phase"]] + context_rows = sum(step["qlens"][: step["contexts"]]) + for name, rows in ( + ("context", slice(0, context_rows)), + ("generation", slice(context_rows, None)), + ): + if output[rows].numel(): + element, rms = _ulp_stats(output[rows], ref[rows]) + print( + f"\n[{cell.why}] {name} rows: max element {element:.2f} ulp, rms {rms:.3f} ulp" + ) + _assert_in_bands(output, ref, bands, context_rows) + + if k is not None: + # The call's own cache write: the live K / V rows land, as a plain E4M3 cast, at their positions. + pool = manager.get_buffers(rt.DENSE, "HND") + first = 0 + for rid, c, n in zip(step["rids"], step["cached"], step["qlens"]): + slots = rt.flat_slots(manager, rt.DENSE, rid, torch.arange(c, c + n)).long() + for role, src in ((0, k), (1, v)): + got = pool[slots // rt.PAGE, role, 0, slots % rt.PAGE] + want = src[first : first + n].to(torch.float8_e4m3fn) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose query rows are shifted by one token is outside it. + shifted = op.reference(backend, _shift(q), k, v, idx_q, metadata, output, sparse, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, shifted, bands, context_rows) + if sparse: + # And the selection is load-bearing: attending every causal block instead is outside it too. + dense = op.reference(backend, q, k, v, idx_q, metadata, output, False, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, dense, bands, context_rows) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py new file mode 100644 index 000000000000..3c5b75f264ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-token valid-block count is written only before the step's first layer. + +The long-row select (``kernels/m3_select_long_dsl.py``, ``N_VALID_PREWAIT``) reads its row's entry of +``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the +buffer happens before the step's first layer: + +- the allocation (``_alloc_msa_proxy_scratch``); +- the host-to-device copy in ``prepare()`` (``_build_step_plans``); +- ``on_update_kv_lens``, which the runner calls before the forward. It writes through the plain + ``<<<>>>`` launch of ``msa_kv_update_kernel`` or through torch copies, so any PDL grid of the step starts only + after it completes. + +This test scans ``tensorrt_llm/_torch`` and fails on any other write site (including writes through a local alias of +the buffer), so a new writer has to revisit the pre-wait read. A buffer handed to an opaque kernel launch is not +tracked; such a launch inside a layer must be reviewed by hand. + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +_NAME = "msa_n_valid_blocks" +# In-place tensor methods that write their receiver. +_MUTATORS = { + "copy_", + "fill_", + "zero_", + "index_put_", + "index_copy_", + "scatter_", + "masked_fill_", + "add_", + "sub_", +} +# Calls whose first argument is the destination they write: ``step_inputs.copy_(dst, src)`` and the H2D helpers. +_COPY_HELPERS = {"copy_", "copy_h2d_if_changed", "copy_h2d_pool_shared"} +_BACKEND = "attention/backends/sparse/minimax_m3/msa_backend.py" +# (file relative to tensorrt_llm/_torch, enclosing function) -> writes allowed there +_ALLOWED = { + (_BACKEND, "_alloc_msa_proxy_scratch"), # the allocation + (_BACKEND, "_build_step_plans"), # the H2D copy in prepare() + (_BACKEND, "on_update_kv_lens"), # the device-side correction before the forward +} + + +def _base(node): + """The innermost receiver of an attribute / subscript / call chain.""" + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute) and node.attr == _NAME: + return node + node = node.func if isinstance(node, ast.Call) else node.value + return node + + +def _mentions(node, aliases) -> bool: + """Whether the expression's receiver chain is ``*.msa_n_valid_blocks`` or a local alias of it.""" + b = _base(node) + if isinstance(b, ast.Attribute): + return b.attr == _NAME + return isinstance(b, ast.Name) and (b.id == _NAME or b.id in aliases) + + +def _refers(node, aliases) -> bool: + """Whether any sub-expression (e.g. a branch of a conditional) names the buffer or an alias.""" + return any( + (isinstance(n, ast.Attribute) and n.attr == _NAME) + or (isinstance(n, ast.Name) and (n.id == _NAME or n.id in aliases)) + for n in ast.walk(node) + ) + + +def _write_sites(tree: ast.AST): + sites = [] + + def visit(node, func, aliases): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + func, aliases = node.name, set() + if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + value = node.value + for t in targets: + if isinstance(t, ast.Name) and value is not None and _refers(value, aliases): + aliases.add(t.id) # a local view of the buffer + elif ( + not isinstance(t, ast.Name) + and _mentions(t, aliases) + and not isinstance(node, ast.AnnAssign) + ): + sites.append((func, node.lineno)) + elif ( + isinstance(t, ast.Attribute) + and t.attr == _NAME + and isinstance(node, ast.AnnAssign) + ): + pass # a dataclass field declaration + if isinstance(node, ast.Call): + f = node.func + if isinstance(f, ast.Attribute) and f.attr in _MUTATORS and _mentions(f.value, aliases): + sites.append((func, node.lineno)) + name = ( + f.id + if isinstance(f, ast.Name) + else (f.attr if isinstance(f, ast.Attribute) else "") + ) + if name in _COPY_HELPERS and node.args and _mentions(node.args[0], aliases): + sites.append((func, node.lineno)) + for child in ast.iter_child_nodes(node): + visit(child, func, aliases) + + visit(tree, None, set()) + return sites + + +def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: + found = set() + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if _NAME not in src: + continue + rel = path.relative_to(_ROOT).as_posix() + for func, line in _write_sites(ast.parse(src)): + assert (rel, func) in _ALLOWED, ( + f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " + "wait (kernels/m3_select_long_dsl.py, N_VALID_PREWAIT), so a writer inside the " + "step must be ordered before the select's grid or the read moved after the wait" + ) + found.add((rel, func)) + assert found == _ALLOWED, f"expected write sites {_ALLOWED}, found {found}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..f759d0d20d3a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_producer entry: the pre-wait-staged NVFP4 QKV + indexer producer. + +Every cell is held to the native producer (``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert``) bit for bit on q, +index-q and every K / V / scale / index-K cache byte, and to the fp32 torch reference: the FP8 outputs within one E4M3 +step, the decoded NVFP4 K / V within one E2M1 step of their block scale. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_producer import ( # noqa: E501 + FP4_BLOCK, + HEADS_Q, + PAGE, + decode_nvfp4_token, + scale_index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_producer import ( # noqa: E501 + m3_producer as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_producer requires a CUDA device" + +PAGES = 1024 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _caches(): + return ( + torch.zeros(PAGES, 2, 1, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, 1, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _inputs(tokens: int, scale: float, bad_slots: bool, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + weights = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + packed = ( + torch.randn(tokens, (HEADS_Q + 4) * 128, device="cuda", generator=g) * scale + ).bfloat16() + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return packed, slots, quant, weights, rotary, positions + + +def _args(inputs, caches, early: bool = True): + packed, slots, quant, weights, rotary, positions = inputs + return (packed, *caches, slots, quant, EPS, *weights, rotary, positions, HEADS_Q, True, early) + + +def _native(inputs, caches): + packed, slots, quant, weights, rotary, positions = inputs + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, *caches, slots, quant, HEADS_Q, 1, 1, 128, 64, EPS, *weights, rotary, positions + ) + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + early = spec.get("early", True) + inputs = _inputs( + spec["tokens"], spec["scale"], spec.get("bad_slots", False), zlib.crc32(cell.why.encode()) + ) + stock_caches, caches = _caches(), _caches() + stock = _native(inputs, stock_caches) + args = _args(inputs, caches, early) + certifying(op) + ref = op.reference(*args) + out = op(*args) + torch.cuda.synchronize() + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, s, o in zip(names, (*stock, *stock_caches), (*out, *caches)): + assert torch.equal(s.view(torch.uint8), o.view(torch.uint8)), ( + f"{name} differs from the native producer" + ) + + _within_fp8(out[0], ref["q"], "q") + _within_fp8(out[1][:, 0], ref["index_q"], "index_q") + kv_data, kv_scale, index_k = caches + slots = inputs[1] + checked = 0 + for t in range(spec["tokens"]): + slot = int(slots[t]) + if slot < 0 or slot >= PAGES * PAGE: + continue + page, within = divmod(slot, PAGE) + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") + for plane, name in ((0, "k"), (1, "v")): + scales = kv_scale[page, plane, 0].reshape(-1)[scale_index(within, plane).cuda()] + dq = decode_nvfp4_token(kv_data[page, plane, 0, within], scales, KV_QUANT[1 + plane]) + r = ref[name][t].bfloat16().float() + # Round to nearest E2M1 under the block's E4M3-rounded scale s (in value units s / quant_scale): at most + # half the widest E2M1 step (one unit), plus the clipping at 6 units when s rounded down; a scale that + # underflows E4M3 zeroes the block. A slack of one E4M3 step covers the kernel's approximate reciprocals. + qs = KV_QUANT[1 + plane] + amax = r.view(-1, FP4_BLOCK).abs().amax(dim=1) + unit = (qs * amax / 6.0).to(torch.float8_e4m3fn).float() / qs + bound = torch.where( + unit > 0, 1.15 * unit + (amax - 6.0 * unit / 1.15).clamp_min(0.0), amax + ) + bound = bound.repeat_interleave(FP4_BLOCK) + 1e-12 + assert bool(((dq - r).abs() <= bound).all()), (name, t) + checked += 1 + assert checked >= spec["tokens"] - 2 + + +def test_early_trigger_does_not_change_the_bits() -> None: + inputs = _inputs(48, 1.0, False, 7) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=True)) + out_b = op(*_args(inputs, b, early=False)) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +@pytest.mark.parametrize( + "tokens, scale, bad_slots, early", + [ + (1, 1.0, False, True), + (4, 3.0, False, True), + (16, 1e-3, False, True), + (32, 40.0, True, True), + (48, 1.0, False, True), + (1000, 1.0, False, False), + ], +) +def test_prewait_step_inputs_is_bitwise_the_shared_op(tokens, scale, bad_slots, early) -> None: + """dispatch.PRODUCER_PREWAIT_STEP_INPUTS: this target's copy (minimax_m3_sm107_tp4_producer_pw, the position ids, + slots and RoPE rows read before its grid wait) writes exactly the shared op's q, index-q and cache bytes.""" + assert hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp4_producer_pw"), ( + "the target's producer op is not loaded (its TP4 extension)" + ) + inputs = _inputs(tokens, scale, bad_slots, 11) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=early)) + out_b = op(*_args(inputs, b, early=early), prewait_step_inputs=True) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +def test_guard_refuses_a_repeated_slot() -> None: + inputs = list(_inputs(4, 1.0, False, 3)) + inputs[1] = inputs[1].clone() + inputs[1][1] = inputs[1][0] + with pytest.raises(AssertionError, match="must not repeat a slot"): + with validating(op): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_operand_off_the_device() -> None: + inputs = list(_inputs(4, 1.0, False, 9)) + inputs[5] = inputs[5].cpu() # position ids on the host + with pytest.raises(RuntimeError, match="packed's CUDA device"): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_unaligned_index_k_page_stride() -> None: + """The index-K rows are written as 32-bit words: a page stride that is not a multiple of 4 is refused.""" + kv_data, kv_scale, _ = _caches() + backing = torch.zeros(PAGES * (128 * 128 + 1), dtype=torch.float8_e4m3fn, device="cuda") + index_k = backing.as_strided((PAGES, 1, PAGE, 128), (128 * 128 + 1, 128 * 128, 128, 1)) + with pytest.raises(RuntimeError, match="index_k"): + op(*_args(_inputs(4, 1.0, False, 13), (kv_data, kv_scale, index_k))) + + +def test_op_refuses_unaligned_weights_or_rope_table() -> None: + """The norm weights are read as 8-byte words and the RoPE rows as float4: offset views are refused.""" + packed, slots, quant, weights, rotary, positions = _inputs(4, 1.0, False, 17) + offset_weight = torch.zeros(129, dtype=torch.bfloat16, device="cuda")[1:] + with pytest.raises(RuntimeError, match="norm weights"): + op( + *_args( + (packed, slots, quant, [offset_weight, *weights[1:]], rotary, positions), _caches() + ) + ) + flat = torch.zeros(rotary.numel() + 1, device="cuda") + offset_rope = flat[1:].view(rotary.shape) + with pytest.raises(RuntimeError, match="rotary_cos_sin"): + op(*_args((packed, slots, quant, weights, offset_rope, positions), _caches())) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..071795f8c3fa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_qk_norm_rope_fp8 entry of the MiniMax-M3 tp4 target: the out-of-place FP8 form of the fused +per-head Gemma RMSNorm + partial NeoX RoPE against the in-place kernel followed by the E4M3 cast. V is a cast, so it +matches bit for bit; Q and K round once from fp32 where the composition rounds through bf16 first, so each element is +within one E4M3 step of it.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 + m3_qk_norm_rope_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" + +HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 + + +def _args(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q, kv = spec["q"], spec["kv"] + qkv = torch.randn(spec["m"], (q + 2 * kv) * HEAD_DIM, generator=g, device="cuda").to( + torch.bfloat16 + ) + qw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (spec["m"],), generator=g, device="cuda", dtype=torch.int32) + return ( + qkv, + q, + kv, + kv, + HEAD_DIM, + ROTARY_DIM, + EPS, + qw, + kw, + THETA, + True, + pos, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = args[0].clone() + with validating(op): + out = op(*args) + assert torch.equal(args[0], before), "the FP8 form must leave qkv untouched" + ref = op.reference(*args) + assert out.dtype == torch.float8_e4m3fn and out.shape == ref.shape + v_start = (cell.spec["q"] + cell.spec["kv"]) * HEAD_DIM + assert torch.equal(out[:, v_start:].view(torch.uint8), ref[:, v_start:].view(torch.uint8)) + assert _e4m3_steps(out[:, :v_start], ref[:, :v_start]) <= 1 + + +def _e4m3_steps(a, b) -> int: + """The largest distance between two E4M3 tensors in representable values (E4M3 codes order by magnitude).""" + + def signed(t): + c = t.view(torch.uint8).int() + # Sign-magnitude to a line: -0 and +0 both map to 0. + return torch.where(c >= 128, 128 - c, c) + + return int((signed(a) - signed(b)).abs().max()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..a7ee9f877707 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_qkv_insert entry: the generic horizontal producer on an FP8 main cache. + +Q, index-Q and every written K / V / index-K cache row are held within one E4M3 step of the fp32 reference. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_qkv_insert import ( # noqa: E501 + m3_qkv_insert as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" + +PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 16, 1, 1, 128, 64 +MAX_POSITIONS = 262144 +EPS = 1e-6 + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +def _args(tokens: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + tokens, (HEADS + 2 * KV + INDEX + 1) * D, device="cuda", generator=g + ).bfloat16() + kv = torch.zeros(PAGES, 2, KV, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + index_k = torch.zeros(PAGES, 1, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + weights = [(torch.randn(D, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + rotary = torch.randn(MAX_POSITIONS, 2, ROTARY // 2, device="cuda", generator=g) + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + return ( + packed, + kv, + None, + index_k, + slots, + None, + HEADS, + KV, + INDEX, + D, + ROTARY, + EPS, + *weights, + rotary, + positions, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _args(cell.spec["m"], zlib.crc32(cell.why.encode())) + certifying(op) + ref = op.reference(*args) + q, index_q = op(*args) + torch.cuda.synchronize() + _within_fp8(q.view(ref["q"].shape), ref["q"], "q") + _within_fp8(index_q.view(ref["index_q"].shape), ref["index_q"], "index_q") + kv, index_k, slots = args[1], args[3], args[4] + for t in range(cell.spec["m"]): + page, within = divmod(int(slots[t]), PAGE) + _within_fp8(kv[page, 0, :, within], ref["k"][t], f"k token {t}") + _within_fp8(kv[page, 1, :, within], ref["v"][t], f"v token {t}") + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..6f23d21ce10f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_blocks entry of the MiniMax-M3 tp4 target: MSA top-16 block selection with the forced +init / local blocks and per-row valid-block masks.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_blocks import ( # noqa: E501 + m3_select_blocks as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" + +TOPK, INIT, LOCAL = 16, 0, 1 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + n_blocks, total_q = spec["n_blocks"], spec["total_q"] + # Distinct scores (a permutation) so the selection has no ties to break. + scores = ( + torch.randperm(n_blocks * total_q, generator=g, device="cuda") + .float() + .view(1, n_blocks, total_q) + ) + if spec["mixed"]: + n_valid = torch.randint( + 17, n_blocks + 1, (total_q,), generator=g, device="cuda", dtype=torch.int32 + ) + else: + n_valid = torch.full((total_q,), n_blocks, dtype=torch.int32, device="cuda") + with validating(op): + out = op(scores, TOPK, n_valid, INIT, LOCAL, True) + ref = op.reference(scores, TOPK, n_valid, INIT, LOCAL, True) + assert out.shape == ref.shape + assert torch.equal(out, ref) + assert out.permute(1, 0, 2).is_contiguous(), "head-major backing" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..6f4d512ba75e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_long entry (top-16 MSA block selection for long decode rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long import ( # noqa: E501 + m3_select_long as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key +# widths, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 2048, 2049, 4096, 4097, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # re-read path). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + # Odd tokens: almost everything -inf, so picks run past the finite scores. + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +def _mass_tie_rows(scores, nvb, init_blocks, local_blocks): + """[total_q, num_kv_heads] mask of rows where more than 1024 valid blocks share the 16th-largest score.""" + num_kv_heads, num_blocks, total_q = scores.shape + blocks = torch.arange(num_blocks, device=scores.device) + raw = nvb[:total_q].long() + valid = raw.clamp(0, num_blocks) + forced = scores.permute(2, 0, 1).clone() # [total_q, heads, blocks] + forced = torch.where(blocks < init_blocks, 1.0e30, forced) + forced = torch.where(blocks >= (raw - local_blocks).clamp_min(0)[:, None, None], 1.0e29, forced) + in_row = (blocks[None, :] < valid[:, None])[:, None, :] + ranked = torch.where(in_row, forced, float("-inf")).sort(dim=-1, descending=True).values + kth = ranked[..., min(TOPK, num_blocks) - 1 : min(TOPK, num_blocks)] + return ((forced == kth) & in_row).sum(dim=-1) > 1024 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + + # Selection is an exact claim: the same blocks or a different attention. + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # And it is the in-tree selector's answer, bit for bit (the op the target falls back to) -- except on rows + # where more than 1024 blocks tie at the 16th-largest score: there the in-tree op does not break the tie to + # the smallest blocks, while this op and the reference do. + scores, nvb, _, init_blocks, local_blocks, _ = args + in_tree = select_blocks_from_maxscore( + scores, + topk=TOPK, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + mass_tie = _mass_tie_rows(scores, nvb, init_blocks, local_blocks) + assert torch.equal(out[~mass_tie], in_tree[~mass_tie]), ( + "selected blocks differ from the in-tree minimax_m3_select_blocks" + ) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + """Releasing the dependent grid right after the select's own wait changes only when the next kernel may launch.""" + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + """Releasing the dependent grid once the row's scores are loaded changes only when the next kernel may launch.""" + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +@pytest.mark.parametrize("cell", range(len(op.CELLS))) +def test_warp_tail_selects_the_same(cell: int) -> None: + """The warp-0 tail (rows of at most 64 candidates) and the all-thread rank pick the same blocks + (dispatch.SELECT_LONG_SL3 and its kill value).""" + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=7 + cell) + rank_out = torch.empty_like(out) + for phase in (0, 1): + op(scores, nvb, out, init_blocks, local_blocks, True, False, phase, True) + op(scores, nvb, rank_out, init_blocks, local_blocks, True, False, phase, False) + assert torch.equal(out, rank_out) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) + + +def test_guard_refuses_out_aliasing_scores() -> None: + scores, nvb, _, init_blocks, local_blocks, _ = _case(op.CELLS[0].spec, seed=3) + aliased = ( + scores.view(torch.int32) + .reshape(-1)[: scores.shape[2] * TOPK] + .view(scores.shape[2], 1, TOPK) + ) + with pytest.raises(AssertionError, match="alias"): + with validating(op): + op(scores, nvb, aliased, init_blocks, local_blocks, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py new file mode 100644 index 000000000000..9866b1d589c0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp4 (Rubin) m3_select_long_dsl entry: the long-row top-16 select from one 384-thread CTA per row, +against the reference and bit for bit against the 1024-thread m3_select_long on every cell.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( + m3_select_long_dsl as select_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, + m3_select_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long_dsl import ( # noqa: E501 + m3_select_long_dsl as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key +# widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 1536, 1537, 4224, 4225, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # whole-row re-rank). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # Every row, mass ties included, is the 1024-thread kernel's answer. + scores, nvb, _, init_blocks, local_blocks, _ = args + long_out = torch.empty_like(out) + m3_select_long(scores, nvb, long_out, init_blocks, local_blocks, False) + assert torch.equal(out, long_out), "selected blocks differ from m3_select_long" + + +@pytest.mark.parametrize("layout", ["kv_blocks_q", "kv_q_blocks"]) +@pytest.mark.parametrize("valid", [17, 1294, 2524, 3907, 9216]) +def test_layouts_and_widths_match_the_1024_thread_select(layout, valid) -> None: + """Both score layouts the scorer can hand over, at a valid count in each key width.""" + g = torch.Generator(device="cuda").manual_seed(valid) + nb = max(valid, 129) + s = torch.randn(1, nb, 8, generator=g, device="cuda") + if layout == "kv_q_blocks": + s = s.permute(0, 2, 1).contiguous().permute(0, 2, 1) + nvb = torch.tensor([valid - (t % 2) for t in range(8)], dtype=torch.int32, device="cuda") + out = torch.empty(1, 8, TOPK, dtype=torch.int32, device="cuda").permute(1, 0, 2) + ref = torch.empty_like(out) + op(s, nvb, out, 0, 1, False) + m3_select_long(s, nvb, ref, 0, 1, False) + assert torch.equal(out, ref) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +@pytest.mark.parametrize("early_trigger", [False, True]) +@pytest.mark.parametrize("cell", range(len(op.CELLS))) +def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: + """kernels/m3_select_long_dsl.N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the grid + wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and tied + rows).""" + phase = 0 if early_trigger else 1 + spec = op.CELLS[cell].spec + scores, nvb, before, init_blocks, local_blocks, _ = _case(spec, seed=11 + cell) + after = torch.empty_like(before) + select_kernels._compile((True, early_trigger, phase, True))( + scores, nvb, before, init_blocks, local_blocks + ) + select_kernels._compile((True, early_trigger, phase, False))( + scores, nvb, after, init_blocks, local_blocks + ) + assert torch.equal(before, after) + + +def test_graph_replay_selects_the_same() -> None: + spec = op.CELLS[1].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=9) + ref = op.reference(scores, nvb, out, init_blocks, local_blocks) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + for _ in range(3): + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..cfd8a14643f6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,442 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode entry (the target's Triton block-sparse decode and its routes).""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_report, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode as stock_decode, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import ( # noqa: E402 — needs the path declared above + HEAD_DIM, + NUM_HEADS, + build_step, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_sparse_decode requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each [16, 128] head row's largest magnitude, and of the relative RMS. Every route +# computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference +# is the BF16 rounding on both sides (one ulp of the row at most), the FP16 probabilities before PV, the split-K +# partials kept in FP32 and merged by their log2-sum-exp, and approximate exp2/log2: 2 ulps per element, the two +# roundings (each at most half an ulp) dominating the RMS, so half an ulp. Measured on GB300 over the cells: at most +# 1.80 ulps per element and 0.20 ulps RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, output, identity): + return op( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + identity, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=spec["identity"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + certifying(op) + _call(case, output, spec["identity"]) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + spec["identity"], + ) + assert torch.isfinite(output.float()).all(), "padding rows must be zeros, not NaN" + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_flag_does_not_change_the_answer() -> None: + """On an identity table the identity routes (the producer, or the multi-query producer from 17 rows) and the + selected-table route (the long-context kernel) compute the same attention.""" + case = build_step((2048, 130, 5, 0, 1777, 900, 2000, 0), 4, identity=True, seed=17) + outs = [] + for identity in (True, False): + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, identity) + outs.append(output) + assert_within_ulp(outs[0], outs[1], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_guard_refuses_a_padded_selection_slot() -> None: + """The Triton producers index the block table with every one of a token's first min(16, blocks) ids, + -1 included.""" + case = build_step((1500,), 1, identity=True, seed=23) + case.topk[0, 0, 3] = -1 + output = torch.empty(1, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="blocks of the request"): + with validating(op): + _call(case, output, True) + + +def test_guard_refuses_a_false_identity_promise() -> None: + """identity_blocks lets the multi-query producer read one table for the whole request.""" + case = build_step((1500, 1600, 1700, 1800, 1500, 1600, 1700, 1800), 4, identity=True, seed=29) + case.topk[0, 5, :2] = case.topk[0, 5, :2].flip(0) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="identity_blocks promises"): + with validating(op): + _call(case, output, True) + + +def _stock(case, output): + stock_decode( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + sm_scale=case.sm_scale, + output=output, + decode_query_len=case.decode_query_len, + k_block_scale=case.k_block_scale, + v_block_scale=case.v_block_scale, + k_global_scale=case.k_global_scale, + v_global_scale=case.v_global_scale, + ) + + +def _identity_table( + seq_lens: torch.Tensor, decode_query_len: int, heads: int, lag: int = 0 +) -> torch.Tensor: + """The identity selection of every verify token, staged from lengths `lag` tokens behind seq_lens.""" + seq = (seq_lens.cpu().to(torch.int64) - lag).clamp(min=0).repeat_interleave(decode_query_len) + intra = torch.arange(seq.shape[0]) % decode_query_len + n_valid = ((seq - decode_query_len + intra + 1).clamp(min=0) + 127) // 128 + ar = torch.arange(16) + table = torch.where(ar[None, :] < n_valid[:, None], ar[None, :], -1).to(torch.int32) + return table.unsqueeze(0).expand(heads, -1, -1).contiguous().cuda() + + +# Each crosses a 128-token edge within its four verify tokens or between the stale and the true lengths. +_EDGE_LENS = (1284, 1287, 1281, 129, 132, 2048, 1028, 5, 643, 1920, 256, 1) + + +@pytest.mark.parametrize("fault", ["stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_identity_route_reads_its_own_lengths_not_the_table(requests, fault, monkeypatch) -> None: + """The Triton identity route (single / pair producer, up to 16 rows) derives block i for slot i from the kv lengths + the kernel reads: a staged table that lags the lengths by a verify step (a 128-token edge crossed), or is garbage, + gives the stock decode's rows over the true identity selection bit for bit. (The LC identity route has its own test + below.)""" + _triton_route(monkeypatch) + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=400 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + assert torch.equal(case.topk, true_table), ( + "build_step's identity table is the identity selection" + ) + if fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _stock(case, stock) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + _call(case_staged, target, True) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("requests", [1, 2, 4]) +def test_identity_route_is_bitwise_the_stock_decode(requests, monkeypatch) -> None: + """The target's Triton producer + split-K merge takes the stock in-tree decode's split-K at every row count + (4 / 8 / 16 identity rows here) and so writes its rows bit for bit.""" + _triton_route(monkeypatch) + case = build_step((1500, 2048, 900, 130)[:requests], 4, identity=True, seed=60 + requests) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _call(case, target, True) + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), case.seq_lens.tolist() + + +def _triton_route(monkeypatch) -> None: + """Route selected steps to the target's Triton producer + merge (this target takes the LC whenever it applies).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( + m3_attn_lc, + ) + + monkeypatch.setattr(m3_attn_lc, "lc_applies", lambda *args, **kwargs: False) + + +def _short_lens(requests: int, g: torch.Generator) -> tuple: + """Short contexts of 1..16 blocks of 128 tokens (the identity-route regime), a few exactly on a block edge.""" + lens = torch.randint(4, 16 * 128 + 1, (requests,), generator=g).tolist() + if requests > 1: + lens[0] = 128 * int(torch.randint(1, 17, (1,), generator=g)) + return tuple(lens) + + +@pytest.mark.parametrize("seed", range(12)) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_short_table_stress_is_bitwise_the_stock_decode(requests, seed, monkeypatch) -> None: + """The fail-closed id / page guard leaves valid inputs alone: short contexts (1..16 blocks) at 1..4 requests x 4 + verify rows, the identity route and a random selected route with -1 tails past each token's blocks, write the + stock decode's rows bit for bit.""" + _triton_route(monkeypatch) + g = torch.Generator().manual_seed(1000 * requests + seed) + for identity in (True, False): + case = build_step(_short_lens(requests, g), 4, identity=identity, seed=7 * seed + requests) + target = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + stock = torch.empty_like(target) + _call(case, target, identity) + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + identity, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("fault", ["neg_in_window", "id_past_row", "page_out_of_pool"]) +def test_invalid_slot_fails_closed(fault, monkeypatch) -> None: + """A corrupt decode input (a -1 inside a token's block window, an id past the block-table row, or a block-table + page outside the pool) is read as an absent block instead of a wild KV read: the rows equal the reference over the + remaining blocks.""" + _triton_route(monkeypatch) + case = build_step((1500, 2048, 700), 4, identity=False, seed=303) + tok, slot = 5, 2 + if fault == "neg_in_window": + case.topk[0, tok, slot] = -1 + elif fault == "id_past_row": + case.topk[0, tok, slot] = case.block_table.shape[1] + 7 + else: + req = tok // 4 + blk = int(case.topk[0, tok, slot]) + case.block_table = case.block_table.clone() + case.block_table[req, blk] = case.k_paged.shape[0] + 1000 + # Every token of the request that selected that block loses it. + hits = (case.topk[:, req * 4 : req * 4 + 4] == blk).nonzero().tolist() + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(case, output, False) + torch.cuda.synchronize() + # The reference walks each token's first min(16, blocks) entries and drops -1: the corrupt slots as -1. + ref_topk = case.topk.clone() + if fault == "page_out_of_pool": + for h, t, k in hits: + ref_topk[h, req * 4 + t, k] = -1 + else: + ref_topk[0, tok, slot] = -1 + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + ref_topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert torch.isfinite(output.float()).all() + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +@pytest.mark.parametrize("fault", ["none", "stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_lc_identity_route_reads_its_own_lengths_not_the_table( + requests, fault, monkeypatch +) -> None: + """An identity step runs the long-context kernel, which takes block i for CTA i from the kv lengths it reads. + Whatever the staged table holds (true, lagging a verify step, all -1, garbage), the rows are the LC's on the true + identity table (the selected-table route) bit for bit, and in the band of the float64 reference.""" + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=700 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + if fault == "none": + staged = true_table.clone() + elif fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + selected = torch.empty_like(target) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + case_true = SimpleNamespace(**{**vars(case), "topk": true_table}) + taken = [] + monkeypatch.setattr(m3_report, "taken", lambda path, rows, what="": taken.append(path)) + _call(case_staged, target, True) + monkeypatch.undo() + # The LC engagement, plus its one-division merge's. + assert sorted(taken) == sorted(["m3_sparse_decode_lc", "sparse_decode_lc_one_div_merge"]), taken + _call(case_true, selected, False) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), selected.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + true_table, + case.block_table, + case.seq_lens, + case.sm_scale, + target, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert_within_ulp(target, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +@pytest.mark.parametrize("requests", [1, 2]) +def test_lc_identity_step_reads_q_after_the_wait(requests, monkeypatch) -> None: + """An identity step has no select, so the LC must read q after its grid wait: the op gets identity_step (and the + device identity selection) on every identity step.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_attn_lc, + m3_prequant, + ) + + lens = tuple(_EDGE_LENS[(requests * 3 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=900 + requests) + true_table = _identity_table(case.seq_lens, 4, case.topk.shape[0]) + case_true = SimpleNamespace(**{**vars(case), "topk": true_table}) + rows = case.q.shape[0] + seen = [] + real = m3_attn_lc.m3_sparse_decode_lc + + def spy(*args, **kwargs): + seen.append((kwargs.get("identity", False), kwargs.get("identity_step", False))) + return real(*args, **kwargs) + + monkeypatch.setattr(m3_attn_lc, "m3_sparse_decode_lc", spy) + out = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(case_true, out, True) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert seen[0] == (True, True), seen + assert bool(torch.isfinite(out.float()).all()) + + +@pytest.mark.parametrize("requests", [2, 3]) +def test_lc_fails_closed_on_ids_past_the_row(requests) -> None: + """A selected id past the token's row (a stale table, e.g. a padded row keeping a 500K request's n_valid) is + skipped, as the Triton producer skips it: the rows equal those with that slot at -1, bit for bit, and stay finite + even though the block-table tail it would reach points at a page whose block scales are NaN (0xFF).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_prequant, + ) + + lens = (165521, 700, 3000)[:requests] + case = build_step(lens, 4, identity=False, seed=1300 + requests) + dql = case.decode_query_len + good = case.topk.clone() + bad = case.topk.clone() + for tok in range(case.q.shape[0]): + req, qoff = divmod(tok, dql) + kv_len = max(int(case.seq_lens[req]) - dql + qoff + 1, 0) + nblk = (kv_len + 127) // 128 + past = nblk + 1 + if 2 <= nblk < 16 and past < case.block_table.shape[1]: + bad[:, tok, nblk - 1] = past + good[:, tok, nblk - 1] = -1 + page = int(case.block_table[req, past]) + case.k_block_scale[page].fill_(0xFF) + case.v_block_scale[page].fill_(0xFF) + assert not torch.equal(bad, good) + outs = [] + for table in (bad, good): + out = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(SimpleNamespace(**{**vars(case), "topk": table}), out, False) + outs.append(out) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert torch.isfinite(outs[0].float()).all() + assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py new file mode 100644 index 000000000000..0c0becea2433 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py @@ -0,0 +1,526 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode_lc entry (one-launch long-context sparse decode).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + WIDTH, + dequant_nvfp4_pages, + mxfp8_handoff_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + m3_sparse_decode_lc as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import build_step # noqa: E402 — needs the path declared above +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_sparse_decode_lc requires a CUDA device" + +#: Hand-off buffers are sized as MERGE_MXFP8 sizes them. +_HANDOFF_ROWS = 128 + +# Band, in BF16 ulps (2^-8) of each output row's largest magnitude, and of the relative RMS. +# +# The reference is float64 attention over the exactly dequantized blocks, rounded once to BF16. The kernel's +# operands are exact too (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so what separates the two is: the final +# BF16 rounding on both sides (up to one ulp of the element, at most one of the row's scale), the softmax +# probabilities rounded to FP16 before PV (2^-11 relative each, averaging down over the block), ex2.approx / +# lg2.approx (~2^-22) and FP32 accumulation order. Hence 2 ulps per element; the RMS is dominated by the two +# roundings (each at most half an ulp, about 0.3 in aggregate), so half an ulp. Measured on GB300 over the cells: +# at most 1.91 ulps per element and 0.25 ulps RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, spec, out, q8, sf): + return op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + spec["linear_softmax"], + False, + True, + ) + + +def _buffers(total_q, mxfp8): + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + if not mxfp8: + return out, None, None + q8 = torch.zeros(_HANDOFF_ROWS, WIDTH, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_HANDOFF_ROWS * (WIDTH // 32), dtype=torch.uint8, device="cuda") + return out, q8, sf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=False, + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + total_q = case.q.shape[0] + out, q8, sf = _buffers(total_q, spec["mxfp8"]) + certifying(op) + _call(case, spec, out, q8, sf) + ref = op.reference( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + spec["linear_softmax"], + False, + True, + ) + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + # Padding rows attend nothing and must be exact zeros, not NaN. + assert torch.isfinite(out.float()).all() + if spec["mxfp8"]: + # The hand-off is a deterministic function of the op's own BF16 rows: exact. + ref_q8, ref_sf = mxfp8_handoff_reference(out, _HANDOFF_ROWS) + assert torch.equal(q8[:total_q], ref_q8), ( + "MXFP8 rows differ from the recipe applied to the BF16 rows" + ) + assert torch.equal(sf, ref_sf), "UE8M0 scales differ from the recipe / swizzle" + + +def test_reference_dequant_matches_the_quantizer() -> None: + """The reference reads the paged bytes through its own layout claims (K scales token-major, V scales in + 4x4 token quads). Hold them to the quantizer's own flat output: every slot was written in order, so slot + ``s`` is page ``s // 128``, token ``s % 128``, and the flat scales are token-major by construction.""" + case = build_step((3000,), 1, identity=False, seed=5) + pages = torch.arange(case.k_paged.shape[0], device="cuda") + views = ((case.k_paged, case.k_block_scale, False), (case.v_paged, case.v_block_scale, True)) + for role, (data, scale, quad) in enumerate(views): + flat_data, flat_scale = torch.ops.trtllm.fp4_quantize( + case.source[role].reshape(case.num_slots, 1, -1).contiguous(), + case.quant[role + 1 : role + 2], + 16, + False, + False, + ) + expected = dequant_nvfp4_pages( + flat_data.view(torch.uint8).reshape(-1, 1, 128, 64), + flat_scale.view(torch.uint8).reshape(-1, 1, 128, 8), + pages, + quad_scales=False, + ) + got = dequant_nvfp4_pages(data, scale, pages, quad_scales=quad) + assert torch.equal(got, expected), ( + f"role {role}: paged dequantization disagrees with the quantizer" + ) + + +def test_the_variants_agree() -> None: + """The two softmax recurrences and the hand-off switch must not change the BF16 rows beyond rounding.""" + spec = dict(seq_lens=(5000, 12345), decode_query_len=4) + case = build_step(spec["seq_lens"], 4, identity=False, seed=11) + outs = [] + for linear in (False, True): + for mxfp8 in (False, True): + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + _call(case, dict(linear_softmax=linear), out, q8, sf) + outs.append(out) + for other in outs[1:]: + assert_within_ulp(other, outs[0], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + torch.testing.assert_close(outs[1], outs[0], rtol=0, atol=0) # hand-off on/off: same rows + + +def test_op_refuses_the_uninstantiated_late_trigger() -> None: + case = build_step((3000,), 1, identity=False, seed=3) + out, q8, sf = _buffers(1, True) + with pytest.raises(RuntimeError, match="early_trigger"): + op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + 1, + case.sm_scale, + True, + False, + False, + ) + + +def test_guard_refuses_a_block_past_the_request() -> None: + """A selected id at or past the request's block count reads a page the token cannot attend; with every + position masked the CTA's softmax is NaN.""" + case = build_step((3000,), 1, identity=False, seed=4) + case.topk[0, 0, 0] = 40 + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="block count"): + with validating(op): + _call(case, dict(linear_softmax=False), out, q8, sf) + + +def test_guard_refuses_strided_page_rows() -> None: + """The kernel walks a page's token rows at a fixed 64-byte pitch; a padded row stride is misread.""" + case = build_step((3000,), 1, identity=False, seed=6) + wide = torch.zeros(*case.k_paged.shape[:3], 80, dtype=torch.uint8, device="cuda")[..., :64] + wide.copy_(case.k_paged) + case.k_paged = wide + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="64-byte rows"): + with validating(op): + _call(case, dict(linear_softmax=False), out, q8, sf) + + +def test_guard_refuses_a_narrow_handoff() -> None: + case = build_step((3000,), 1, identity=False, seed=7) + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="num_sf_k_tiles"): + with validating(op): + op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + 8, + 1, + case.sm_scale, + False, + False, + True, + ) + + +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + [ + ((165521,), 4), + ((5000, 12345), 4), + ((700,), 1), + ((1500, 300, 90000), 2), + ((200000, 3000, 64, 70000), 4), + ], + ids=["1x165k", "2req", "short-1row", "3req-mixed", "4req-incl-64"], +) +def test_merge_once_is_bitwise_the_image_op(seq_lens, decode_query_len) -> None: + """dispatch.SPARSE_DECODE_LC_MERGE_ONCE: the target's op (merge weights formed once per CTA) writes exactly the + image op's rows and hand-off, at both softmax recurrences, including requests with fewer than 16 selected blocks + (empty chunks, -inf LSE). With SPARSE_DECODE_LC_PREFETCH_OPROJ it also pulls two ranges the size of the TP4 + o_proj weight and scales (one not a multiple of 16 bytes) into L2: the same bits again. With + SPARSE_DECODE_LC_PREFETCH_NEXT_KV it records each row's selection (the first min(16, blocks) ids, -1 after) and + prefetches a next layer's recorded blocks, stale ids past a request's blocks included: the same bits, and the + record.""" + assert hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp4_sparse_decode_lc"), ( + "the target's sparse decode op is not loaded (its TP4 extension)" + ) + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + ) + o_proj_weight = torch.empty(6144, 2048, dtype=torch.uint8, device="cuda") + o_proj_scale = torch.empty(6144 * 64 + 9, dtype=torch.uint8, device="cuda") + rows = case.q.shape[0] + n1_self = torch.full((128, 16), -7, dtype=torch.int32, device="cuda") + n1_next = torch.randint( + -1, 2 * (max(seq_lens) // 128 + 1), (128, 16), dtype=torch.int32, device="cuda" + ) + next_kv = (n1_self, n1_next, case.k_paged, case.v_paged, case.k_block_scale, case.v_block_scale) + variants = ( + (False, None, None), + (True, None, None), + (True, (o_proj_weight, o_proj_scale), None), + (True, (o_proj_weight, o_proj_scale), next_kv), + ) + for linear in (False, True): + for mxfp8 in (False, True): + got = [] + for merge_once, prefetch, n1 in variants: + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + linear, + False, + True, + merge_once=merge_once, + prefetch=prefetch, + next_kv=n1, + ) + got.append((out, q8, sf)) + torch.cuda.synchronize() + a, aq, asf = got[0] + for i, (b, bq, bsf) in enumerate(got[1:], 1): + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), (linear, mxfp8, i) + if mxfp8: + assert torch.equal(aq, bq) and torch.equal(asf, bsf), (linear, mxfp8, i) + # The record: each row's first min(16, blocks) selected ids, -1 past them; rows past the step untouched. + want = case.topk[0].clone() + for tok in range(rows): + req, qoff = divmod(tok, case.decode_query_len) + kv_len = max(int(case.seq_lens[req]) - case.decode_query_len + qoff + 1, 0) + want[tok, min(16, (kv_len + 127) // 128) :] = -1 + assert torch.equal(n1_self[:rows], want), "the recorded selection" + assert bool((n1_self[rows:] == -7).all()), "rows past the step" + + +def _tp4_op(case, out, q8, sf, linear, use_pdl, one_div): + torch.ops.trtllm.minimax_m3_sm107_tp4_sparse_decode_lc( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + linear, + use_pdl, + True, + one_div=one_div, + ) + + +_ONE_DIV_STEPS = [ + ((165521,), 4), + ((5000, 12345), 4), + ((700,), 1), + ((1500, 300, 90000), 2), + ((200000, 3000, 64, 70000), 4), +] + + +@pytest.mark.parametrize("linear", [False, True]) +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + _ONE_DIV_STEPS, + ids=["1x165k", "2req", "short-1row", "3req-mixed", "4req-incl-64"], +) +def test_one_div_merge_is_deterministic_and_in_band(seq_lens, decode_query_len, linear) -> None: + """The one-division merge (on in this target): it writes the same rows on every run + (5 eager runs, PDL on and off) and under CUDA-graph replay (3 replays == eager), MXFP8 hand-off included; the rows + are within the float64 reference's band, and within summation-order distance of the per-chunk merge.""" + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + 1 + ) + total_q = case.q.shape[0] + runs = [] + for i in range(6): + out, q8, sf = _buffers(total_q, True) + _tp4_op(case, out, q8, sf, linear, i == 5, True) + runs.append((out, q8, sf)) + torch.cuda.synchronize() + a, aq, asf = runs[0] + for b, bq, bsf in runs[1:]: + assert torch.equal(a.view(torch.int16), b.view(torch.int16)) + assert torch.equal(aq, bq) and torch.equal(asf, bsf) + out = torch.empty_like(a) + _tp4_op(case, out, None, None, linear, True, True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _tp4_op(case, out, None, None, linear, True, True) + for _ in range(3): + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), a.view(torch.int16)) + per, _, _ = _buffers(total_q, False) + _tp4_op(case, per, None, None, linear, False, False) + ref = op.reference( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + a, + None, + None, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + linear, + False, + True, + ) + assert torch.isfinite(a.float()).all() + assert_within_ulp(a, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + assert_within_ulp(a, per, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +_BT_PREWAIT_STEPS = [ + ((165521,), 4), + ((5000, 12345), 4), + ((700,), 1), + ((1500, 300, 90000), 2), + ((200000, 3000, 64, 70000), 4), + ((524288,), 4), # exactly the 4096-entry staging cap + ((524289, 2000), 4), # one block past the cap: the global lookup, beside a staged request + ((600000,), 1), +] + + +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + _BT_PREWAIT_STEPS, + ids=[ + "1x165k", + "2req", + "short-1row", + "3req-mixed", + "4req-incl-64", + "at-cap", + "past-cap+short", + "600k-1row", + ], +) +def test_bt_prewait_is_bitwise(seq_lens, decode_query_len) -> None: + """The op's ``bt_prewait`` option (not used by this target): staging the request's block-table row in shared memory + before the grid wait writes exactly the rows and MXFP8 hand-off of the global lookup, at both softmax recurrences, + with and without the one-division merge, PDL on and off, eager and under CUDA-graph replay; rows at the 4096-entry + cap are staged, longer ones keep the global lookup; the o_proj / next-KV prefetch hints alongside change nothing.""" + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + 7 + ) + total_q = case.q.shape[0] + n1_self = torch.full((128, 16), -7, dtype=torch.int32, device="cuda") + n1_next = torch.randint( + -1, 2 * (max(seq_lens) // 128 + 1), (128, 16), dtype=torch.int32, device="cuda" + ) + next_kv = (n1_self, n1_next, case.k_paged, case.v_paged, case.k_block_scale, case.v_block_scale) + + def run(out, q8, sf, linear, one_div, bt, use_pdl, hints=False): + torch.ops.trtllm.minimax_m3_sm107_tp4_sparse_decode_lc( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + linear, + use_pdl, + True, + *((None, None) + next_kv if hints else ()), + one_div=one_div, + bt_prewait=bt, + ) + + for linear in (False, True): + for one_div in (False, True): + base = _buffers(total_q, True) + run(*base, linear, one_div, False, False) + arms = [] + for bt, use_pdl, hints in ( + (True, False, False), + (True, True, False), + (True, True, True), + ): + bufs = _buffers(total_q, True) + run(*bufs, linear, one_div, bt, use_pdl, hints) + arms.append(bufs) + torch.cuda.synchronize() + a, aq, asf = base + for i, (b, bq, bsf) in enumerate(arms): + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), (linear, one_div, i) + assert torch.equal(aq, bq) and torch.equal(asf, bsf), (linear, one_div, i) + out = torch.empty_like(a) + run(out, None, None, linear, one_div, True, True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(out, None, None, linear, one_div, True, True) + for _ in range(3): + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), a.view(torch.int16)), ( + linear, + one_div, + "graph", + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py new file mode 100644 index 000000000000..35f04a37eb8e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for kernels/m3_tokmap_guard: a one-request multi-token call of +thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a +stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the +appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for +every length. Without the guard the stale map moves rows (the negative control), and a several-request call passes +through with its own map untouched. + +The fresh-map reference forces the op's own build (a scaled RoPE type with no inv-freq cache, at rotary dim 0, so no +row is rotated either way). +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_tokmap_guard as guard, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = 16, 1, 128, 128 +BUCKET = 4 +_PEER_EXTRA = 777 # the second request's extra length behind the first +FP8_KV_QUANT_MODE = 128 +_DEV = torch.device("cuda") # unindexed on purpose: install() and the lookup normalize it +_TAKEN = [] + + +@pytest.fixture(scope="module", autouse=True) +def _installed(): + guard.install(_DEV, 48, _TAKEN.append) + yield + + +class _Step: + """One request (or ``nreq``) of BUCKET padded rows over a fresh FP8 paged cache, persistent across calls like a + real step's buffers.""" + + def __init__(self, kv_len: int, nreq: int = 1, seed: int = 0): + g = torch.Generator(device=_DEV).manual_seed(seed + kv_len) + self.nreq = nreq + # Per-request inputs (page tables, generation lengths, position offsets) for at least two requests, the page + # tables at the peer's longer length (_seq_with_peer): the stale map sends rows to request 1, which must have + # real entries behind it, also on the unguarded calls. + tables = max(nreq, 2) + pages = math.ceil((kv_len + _PEER_EXTRA + BUCKET) / PAGE) + self.nblocks = tables * pages + 2 + self.pool0 = ( + torch.randn(self.nblocks, 2, KVH, PAGE, D, device=_DEV, generator=g) + .clamp_(-4, 4) + .to(torch.float8_e4m3fn) + ) + self.pool = self.pool0.clone() + perm = torch.randperm(self.nblocks, device=_DEV, generator=g)[: tables * pages].view( + tables, pages + ) + self.offsets = torch.zeros(1, tables, 2, pages, dtype=torch.int32, device=_DEV) + self.offsets[0, :, 0] = 2 * perm + self.offsets[0, :, 1] = 2 * perm + 1 + self.pointers = torch.tensor([[self.pool.data_ptr(), 0]], dtype=torch.int64) + self.mapping = torch.zeros(1, 2, dtype=torch.int32) + self.qkv = torch.randn( + nreq * BUCKET, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=_DEV, generator=g + ) + # Request lengths after this step's append; generation lengths (accepted rows) set per case. + self.seq = torch.full((nreq,), kv_len, dtype=torch.int32, device=_DEV) + self.gen = torch.full((tables,), BUCKET, dtype=torch.int32, device=_DEV) + self.pos = ( + torch.arange(BUCKET, dtype=torch.int32, device=_DEV).repeat(tables, 1).contiguous() + ) + self.one = torch.ones(1, dtype=torch.float32, device=_DEV) + self.ws = torch.empty(64 << 20, dtype=torch.uint8, device=_DEV) + self.kv_len = kv_len + + def call(self, op, fresh_map: bool = False): + """One generation preprocess; returns a copy of the processed query.""" + out = op( + self.qkv, + self.ws, + self.seq, + self.gen, + self.pos, + self.offsets, + self.pointers, + self.mapping, + self.one, + self.one, + None, + None, + None, + None, + 0, + 0, + H, + KVH, + D, + PAGE, + FP8_KV_QUANT_MODE, + 1 << 20, + 1 << 20, + self.nreq * BUCKET, + self.nreq, + BUCKET, + self.kv_len, + 0, + 10000.0, + 1 if fresh_map else 0, + 1.0, + 1 << 20, + 0, + 1.0, + 1.0, + False, + BUCKET, + 0, + torch.cuda.get_device_properties(0).multi_processor_count, + self.nblocks, + 2, + True, + False, + ) + return out[0].clone() + + def poison_map(self) -> None: + """A stale multi-request map in the op's one-request map slice: rows 1.. on request 1.""" + layout = thop.get_trtllm_gen_generation_workspace_layout( + self.qkv.dtype, 1, self.nreq * BUCKET, H, D, 0, KVH, 0, False, False + ) + off = int(layout["tokens_info_offset"]) + stale = torch.tensor( + [[0, 0]] + [[1, t] for t in range(BUCKET - 1)], dtype=torch.int32, device=_DEV + ) + self.ws[off : off + 8 * BUCKET].view(torch.int32).view(BUCKET, 2).copy_(stale) + + def reset(self) -> None: + self.pool.copy_(self.pool0) + + +def _seq_with_peer(step: _Step) -> None: + # A second request's length behind the first, so a row the stale map moves to request 1 reads a real length. + step.seq = torch.tensor( + [step.kv_len, step.kv_len + _PEER_EXTRA], dtype=torch.int32, device=_DEV + ) + + +@pytest.mark.parametrize("kv_len", [130, 8191, 34117]) +@pytest.mark.parametrize("n", [1, 2, 3, 4]) +def test_guard_equals_a_fresh_map(kv_len, n) -> None: + step = _Step(kv_len) + _seq_with_peer(step) + step.gen.fill_(n) + step.poison_map() + step.reset() + q_ref = step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + + step.poison_map() + step.reset() + q = step.call(thop.trtllm_gen_generation_preprocess) + # The processed query's n live rows (rows past the request's length are padding the op leaves unwritten). + assert torch.equal(q[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)) + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + assert BUCKET in _TAKEN + + +def test_stale_map_moves_rows_without_the_guard() -> None: + step = _Step(34117) + _seq_with_peer(step) + step.gen.fill_(BUCKET) + step.poison_map() + step.reset() + step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + step.poison_map() + step.reset() + step.call(guard._original) + assert not torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + + +def test_one_captured_graph_serves_every_length() -> None: + step = _Step(34000) + _seq_with_peer(step) + refs = {} + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + q = step.call(guard._original, fresh_map=True) + refs[n] = (q, step.pool.clone()) + step.gen.fill_(BUCKET) + step.call(thop.trtllm_gen_generation_preprocess) # eager first, as the warmup runs it + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = thop.trtllm_gen_generation_preprocess + q_static = step.call(out) + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + g.replay() + torch.cuda.synchronize() + q_ref, pool_ref = refs[n] + assert torch.equal(q_static[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)), n + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)), n + + +def test_several_requests_pass_through() -> None: + step = _Step(8191, nreq=2) + layout = thop.get_trtllm_gen_generation_workspace_layout( + step.qkv.dtype, 2, 2 * BUCKET, H, D, 0, KVH, 0, False, False + ) + step.reset() + step.call(thop.trtllm_gen_generation_preprocess) + pool_guarded = step.pool.clone() + step.reset() + step.call(guard._original) + assert torch.equal(step.pool.view(torch.uint8), pool_guarded.view(torch.uint8)) + assert int(layout["tokens_info_size"]) >= 8 * 2 * BUCKET diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..85bb1fd0bcde --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the msa_kv_update entry (the one-launch MSA KV-length correction).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.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_107__tp4.ops.attention.msa_kv_update import ( # noqa: E501 + TOKENS_PER_BLOCK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" + +#: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. +_SPARE = 3 + + +def _case(spec, seed): + """The metadata's buffers for one step, with stale values in every output so a skipped write shows.""" + g = torch.Generator(device="cpu").manual_seed(seed) + batch, qlen = spec["batch"], spec["query_len"] + total_q = batch * qlen + staged = torch.randint(qlen, 5000, (batch,), generator=g, dtype=torch.int32) + accept = spec["accept"] + if accept == "full": + kv = staged.clone() + elif accept == "page_edge": + staged = torch.tensor([130, 257, 2 * TOKENS_PER_BLOCK + qlen][:batch], dtype=torch.int32) + kv = staged - torch.tensor([0, 1, qlen - 1][:batch], dtype=torch.int32) + else: + kv = staged - torch.randint(0, qlen, (batch,), generator=g, dtype=torch.int32) + if accept == "padded": + kv[batch // 2 :] = 0 + staged[batch // 2 :] = 0 + ncols = (5000 + qlen) // TOKENS_PER_BLOCK + 2 + table = ( + torch.randperm(batch * ncols + 7, generator=g)[: batch * ncols] + .to(torch.int32) + .reshape(batch, ncols) + ) + dev = dict(device="cuda", dtype=torch.int32) + stale = -12345 + return ( + torch.cat([kv, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.cat([staged, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.full((batch + _SPARE,), qlen, **dev), + torch.arange(batch, **dev).repeat_interleave(qlen), + torch.arange(qlen, **dev).repeat(batch), + table.cuda(), + TOKENS_PER_BLOCK, + batch, + total_q, + torch.full((batch + _SPARE,), stale, **dev), + torch.full((batch + 1 + _SPARE,), stale, **dev) if spec["cu_kv"] else None, + torch.full((total_q + _SPARE,), stale, **dev), + torch.full((total_q + _SPARE,), stale, **dev) if spec["n_valid"] else None, + ) + + +def _check(args, ref): + seq, cu_kv, loc, n_valid = args[9], args[10], args[11], args[12] + batch, total_q = args[7], args[8] + ref_seq, ref_cu, ref_loc, ref_nv = ref + # Integer bookkeeping: exact, and nothing past the live rows is touched. + assert torch.equal(seq[:batch], ref_seq) + assert torch.equal(loc[:total_q], ref_loc) + assert bool((seq[batch:] == -12345).all()) and bool((loc[total_q:] == -12345).all()) + assert (cu_kv is None) == (ref_cu is None) and (n_valid is None) == (ref_nv is None) + if cu_kv is not None: + assert torch.equal(cu_kv[: batch + 1], ref_cu) + assert bool((cu_kv[batch + 1 :] == -12345).all()) + if n_valid is not None: + assert torch.equal(n_valid[:total_q], ref_nv) + assert bool((n_valid[total_q:] == -12345).all()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + op(*args) + _check(args, op.reference(*args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_hook_takes_every_cell(cell) -> None: + """The target's hook (what the metadata calls) takes each certified step and writes what the op does.""" + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + 1) + assert hook(*args) is True + _check(args, op.reference(*args)) + + +def test_the_hook_declines_what_the_op_would_misread() -> None: + """A strided operand is read wrongly by the op, so the hook declines it (the metadata then runs its own + torch update) and writes nothing.""" + args = list(_case(op.CELLS[2].spec, seed=5)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + before = args[11].clone() + assert hook(*args) is False + assert torch.equal(args[11], before) + + +def test_guard_refuses_a_strided_operand() -> None: + args = list(_case(op.CELLS[2].spec, seed=6)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_index_past_the_batch() -> None: + args = list(_case(op.CELLS[2].spec, seed=7)) + args[3] = args[3].clone() + args[3][-1] = args[7] + with pytest.raises(AssertionError, match="q_batch_row"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py new file mode 100644 index 000000000000..127859c5cb2a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the ``m3_allreduce`` entry, 4 ranks: its cells against the entry's reference, and the +fused draft loop's layer boundaries built on it (``M3Eagle3Draft.ar_norm``: the plain all-reduce, then FlashInfer's +fused add + RMSNorm) against stock Eagle3's, bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy the single-node drafter can run (AUTO, +ONESHOT, TWOSHOT, NCCL; MNNVL is multi-node only), at tp4 shapes (H = 6144, eps = 1e-6, 1 to 32 rows). The inputs are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own 4-rank launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_allreduce.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py +""" + +import math +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 4, 8, 32) +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +RANK = 0 +WORLD = 1 + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on rank 0, x/4 on ranks 1 and 2, 0 on rank 3. + parts = [x / 2, x / 4, x / 4] + [torch.zeros_like(x)] * (WORLD - 3) + return ( + [p.bfloat16() for p in parts[:WORLD]], + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A, -A, +B, -B + d) onto a unit residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a, -a, b, (-b + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return [p.bfloat16() for p in parts[:WORLD]], residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +def _ulps_at(out: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the magnitude ``scale`` (the operands'), not at the result's.""" + return ((out.double() - ref.double()).abs() / _ulp_at(scale.abs())).max().item() + + +def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary (``M3Eagle3Draft.ar_norm``): ``m3_allreduce`` plain, then the catalog's + ``flashinfer_fused_add_rmsnorm``.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + + return flashinfer_fused_add_rmsnorm(m3_allreduce(x, ar, None), residual.clone(), w, EPS) + + +def _check_cells(mapping: Any) -> None: + """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum + to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude + u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final + roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). + The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm.functional import AllReduceStrategy + + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating + + ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + for i, cell in enumerate(m3_allreduce.CELLS): + spec = cell.spec + xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) + params = ( + None + if spec.get("plain") + else AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual.clone(), + norm_weight=w, + eps=EPS, + trigger_completion_at_end=False, + ) + ) + # The reference first: the fused path may write the updated residual into ``params.residual``. + ref = m3_allreduce.reference(xs, ar, params) + with validating(m3_allreduce): + out = m3_allreduce(xs[RANK], ar, params) + torch.cuda.synchronize() + total = sum(x.double() for x in xs) + if params is None: + err = _ulps_at(out, ref, total) + assert err <= 1.0, (cell.why, err) + else: + err = _ulps_at(out[1], ref[1], torch.maximum(total.abs(), residual.double().abs())) + assert err <= 2.5, (cell.why, "residual", err) + h = out[1].double() + own = h * torch.rsqrt(h.square().mean(-1, keepdim=True) + EPS) * w.double() + assert _ulps(out[0], own) <= 2.0, (cell.why, "normed", _ulps(out[0], own)) + dist.barrier() + if RANK == 0: + print(f"{RESULT_TAG} cell {cell.why} passed", flush=True) + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + _check_cells(mapping) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + ulp, res_ulp = 0.0, 0.0 + for m in TOKENS: + xs, residual, w = build(m) + target_out, target_res = _target(ar, xs[RANK], residual, w) + stock_out, stock_res = _stock(ar, xs[RANK], residual, w) + torch.cuda.synchronize() + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {name} m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {name:34s} | target = stock vs fp64: normed {ulp:.3f}" + f" | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the drafter runs at tp4, got {WORLD} ranks" + try: + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py new file mode 100644 index 000000000000..80fa9787751d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CAPTURE_NORM_IN_BOUNDARY on the tp4 multicast all-reduce + RMSNorm op, at 4 ranks. + +With ``cap_out`` the boundary also writes the Eagle3 capture of its updated residual. Asserted bit for bit against +``m3_spec_tail.capture_norm`` of the op's own updated residual, in a strided [M, H] block of a [rows, 3H] capture +buffer (the other bytes untouched), and the op's normed / updated / MXFP8 outputs bitwise the same call's without +``cap_out``. Every kernel a boundary can run: + +- the 8-CTA cluster kernel in the stock oneshot order (options 39 / 47, plain calls up to 8 rows, x and x + x2, + prepushed 2); +- the row kernel (plain options 7 / 15 and the carveout instance 71 / 79, 1 to 8 rows on the 8-token bucket, 12 rows + on the 16-token bucket, 32 rows on the 48-token bucket, prepushed 2); + +(The tp4 MoE finalize runs the shared GB300 op, which has no capture epilogue.) + +at residual scales 1e-3 to 1e3 and the fc_norm eps 1e-6 / 1e-5. + +Its own 4-rank launcher, as ``_m3_draft_ar_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_capture_in_boundary.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_capture_in_boundary_op_matrix.py +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +ar_module: Any = None +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + handle = _make_mnnvl_mcast_buffer(dist.group.WORLD, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=dist.group.WORLD, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor) -> None: + """Store this rank's rows into every rank's exchange (set A), as a producer does before a prepushed call.""" + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = parity.view(m, CLUSTER, 1) * half + (t * WORLD + RANK) * (HIDDEN // 2) + c * chunk + w + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, (where, out.dtype, ref.dtype) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +def _capture_reference(updated, weight, eps, out) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_spec_tail, + ) + + m3_spec_tail.capture_norm(updated, None, weight, eps, out) + + +def _one_case( + kind: str, m: int, maxt: int, opts: int, scale: float, quant: bool, seed: int +) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written, + ) + + cap_eps = (1e-6, 1e-5)[seed % 2] + ws = WORKSPACES[maxt] + residual = (_rows(seed + 7, -1, m).float() * 4.0 * scale).to(torch.bfloat16) + norm_w, cap_w = _norm_w(seed), _norm_w(seed + 101) + results = [] + for with_cap in (False, True): + buf = torch.full((m + 3, 3 * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + cap = buf[:m, HIDDEN : 2 * HIDDEN] if with_cap else None + cw = cap_w if with_cap else None + flags = ws.flags_for(CLUSTER) + x = _rows(seed, RANK, m, scale) + x2 = _rows(seed + 3, RANK, m, scale) if kind == "x2" else None + prepushed = 2 if kind == "prepushed" else 0 + if prepushed: + _push(ws, x) + outs = allreduce_entry.m3_mc_allreduce_norm( + x, x2, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), flags, RANK, WORLD, CLUSTER, False, maxt, + quant, prepushed, opts, cap, cw, cap_eps, + ) # fmt: skip + torch.cuda.synchronize() + results.append((outs, buf)) + dist.barrier() + (plain, _), (capped, buf) = results + where = f"{kind} opts {opts} bucket {maxt} rows {m} scale {scale} quant {quant}" + assert len(plain) == len(capped), where + for i, (a, b) in enumerate(zip(plain, capped)): + if i == 3: + written = sf_written(m, HIDDEN // 32, a.device) + a, b = a[written], b[written] + _assert_bitwise(b, a, f"{where}: output {i} with vs without cap_out") + ref = torch.empty((m, HIDDEN), dtype=torch.bfloat16, device="cuda") + _capture_reference(capped[1], cap_w, cap_eps, ref) + torch.cuda.synchronize() + _assert_bitwise(buf[:m, HIDDEN : 2 * HIDDEN], ref, f"{where}: capture vs capture_norm") + untouched = buf.clone() + untouched[:m, HIDDEN : 2 * HIDDEN] = 7.0 + assert bool((untouched == 7.0).all()), f"{where}: bytes outside the capture block changed" + dist.barrier() + + +def _cases(): + """(kind, rows, bucket, opts, scale, quant).""" + cases = [] + scales = (1.0, 1e3, 1e-3) + # The cluster kernel in the oneshot order (plain calls up to 8 rows): options 39 and 47. + for i, m in enumerate((1, 4, 8)): + for opts in (39, 47): + cases.append(("x" if i != 1 else "x2", m, 8, opts, scales[i], opts == 47)) + cases.append(("prepushed", 4, 8, 39, 1.0, False)) + # The row kernel: plain options 7 / 15 and the carveout instance (+ 64). + for i, (m, opts) in enumerate(((1, 7), (4, 15), (8, 7 | 64), (3, 15 | 64))): + cases.append(("x2" if i % 2 else "x", m, 8, opts, scales[i % 3], i == 1)) + cases.append(("prepushed", 8, 8, 7, 1e-3, False)) + cases.append(("x", 12, 16, 15, 1.0, True)) + cases.append(("x2", 12, 16, 15 | 64, 1e3, False)) + cases.append(("x", 32, 48, 15, 1e-3, False)) + return cases + + +def _run_matrix() -> None: + for n, (kind, m, maxt, opts, scale, quant) in enumerate(_cases()): + _one_case(kind, m, maxt, opts, scale, quant, 7000 + 37 * n) + if RANK == 0: + print( + f"[case] {kind} opts {opts} bucket {maxt} rows {m} scale {scale}: bitwise", + flush=True, + ) + + +def _run_one_rank() -> int: + global RANK, WORLD, allreduce_entry, ar_module + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the tp4 boundary runs at 4 ranks, got {WORLD}" + try: + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + + allreduce_entry, ar_module = _plain, _ar + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_draft_ar_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py new file mode 100644 index 000000000000..7e241a8c484e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py @@ -0,0 +1,533 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused draft loop's layer boundaries (``modeling.M3Eagle3Worker._m3_draft_ar_norm``) against stock Eagle3's, +bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy the single-node drafter can run (AUTO, +ONESHOT, TWOSHOT, NCCL; MNNVL is multi-node only), at tp4 shapes (H = 6144, eps = 1e-6, 1 to 32 rows). The inputs are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +The worker method runs with stand-ins for what it reads (the worker's ``mapping``; the drafter's ``w[norm_key]``, +``rms_eps``, ``allreduce`` and its own ``M3Eagle3Draft.ar_norm``), over the multicast workspace the target's +``m3_mc_allreduce.get_workspace`` returns (built here over the torch.distributed group, as the sibling multicast +matrices build it, and registered where ``get_workspace`` finds it). Two paths per strategy, case and row count: + +- ``boundary``: the method as the target runs it. At the rows ``DRAFT_MC_AR`` covers (1..``DRAFT_MC_AR_MAX_ROWS``) + the multicast all-reduce + residual + RMSNorm in its drafter form (row kernel, opts bit 16; the strategy is not + used); above them the built-in path, ``M3Eagle3Draft.ar_norm`` on the strategy's all-reduce. +- ``declined``: the method with ``get_workspace`` returning None (no multicast workspace), so the covered rows take + its fallback (``missed``, then ``M3Eagle3Draft.ar_norm`` on the strategy's all-reduce) as well. + +Which of the two ran is checked per call (the drafter-form multicast call returned rows exactly at the covered rows +of ``boundary``, never in ``declined``). Not driven: ``pushed=True`` (``DRAFT_AR_PUSH``, prepushed 2; it needs the +projection's split-K push) and ``DRAFT_MC_AR_CLUSTER`` (off; exact-class, not bitwise). + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own 4-rank launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_draft_ar_norm.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_draft_ar_norm_op_matrix.py +""" + +import contextlib +import math +import os +import signal +import socket +import subprocess +import sys +import time +import types +from pathlib import Path +from typing import Any, Callable, Iterator, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 2, 3, 4, 5, 6, 7, 8, 9, 32) +"""Every row count of the drafter-form multicast boundary (1..8: its option table differs by row), the first one +past it and a wide one (the built-in path).""" +NORM_KEYS = ("post_norm", "final_norm") +"""The drafter's two boundaries per pass (o_proj -> post-attention norm, down -> final norm), alternated by row +count: the method reads only ``w[norm_key]``, so they differ in the weight's key alone.""" +CLUSTER = 8 +EMPTY_WORD = -2147483648 +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +RANK = 0 +WORLD = 1 + +# Bound inside the rank body: importing the target pulls in tensorrt_llm. +modeling: Any = None +ar_module: Any = None + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on rank 0, x/4 on ranks 1 and 2, 0 on rank 3. + parts = [x / 2, x / 4, x / 4] + [torch.zeros_like(x)] * (WORLD - 3) + return ( + [p.bfloat16() for p in parts[:WORLD]], + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A, -A, +B, -B + d) onto a unit residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a, -a, b, (-b + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return [p.bfloat16() for p in parts[:WORLD]], residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +class _McCalls: + """Wraps ``m3_mc_allreduce.m3_mc_allreduce_norm`` (the method reads it off the module at each call): counts the + drafter-form calls that returned rows, i.e. the boundaries the multicast kernel ran.""" + + def __init__(self, fn: Callable) -> None: + self.fn = fn + self.ran = 0 + + def __call__(self, *args, **kwargs): + out = self.fn(*args, **kwargs) + if out is not None and kwargs.get("drafter_norm"): + self.ran += 1 + return out + + +@contextlib.contextmanager +def _no_workspace() -> Iterator[None]: + """``get_workspace`` returning None, as the target sees it with no multicast workspace: the method's decline.""" + real = ar_module.get_workspace + ar_module.get_workspace = lambda mapping, hidden: None + try: + yield + finally: + ar_module.get_workspace = real + + +def _stand_ins(mapping: Any, ar: Any, w: torch.Tensor) -> Tuple[Any, Any]: + """(worker, drafter) with what ``M3Eagle3Worker._m3_draft_ar_norm`` reads: the worker's ``mapping``; the + drafter's norm weights, ``rms_eps``, its all-reduce and its own ``ar_norm`` (the target's method, bound).""" + worker = types.SimpleNamespace(mapping=mapping) + draft = types.SimpleNamespace( + w={key: w for key in NORM_KEYS}, rms_eps=EPS, allreduce=ar, dtype=torch.bfloat16 + ) + draft.ar_norm = types.MethodType(modeling.M3Eagle3Draft.ar_norm, draft) + return worker, draft + + +def _target( + mapping: Any, + ar: Any, + x: torch.Tensor, + residual: torch.Tensor, + w: torch.Tensor, + norm_key: str, + mc: _McCalls, + declined: bool, +) -> Tuple[torch.Tensor, torch.Tensor, bool]: + """The draft loop's boundary as the target runs it (``M3Eagle3Worker._m3_draft_ar_norm``, ``pushed`` False); + with ``declined`` the multicast workspace is unavailable. Returns (normed, residual sum, whether the multicast + drafter form ran).""" + worker, draft = _stand_ins(mapping, ar, w) + before = mc.ran + with _no_workspace() if declined else contextlib.nullcontext(): + out, res = modeling.M3Eagle3Worker._m3_draft_ar_norm( + worker, draft, x, residual.clone(), norm_key + ) + return out, res, mc.ran > before + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _register_workspace(mapping: Any) -> Any: + """The multicast workspace ``m3_mc_allreduce.get_workspace(mapping, HIDDEN)`` returns: one bucket per + ``BUCKETS`` entry, each built as ``McAllReduceWorkspace.create`` builds it but over the torch.distributed group + (these ranks have no MPI), with the 8-CTA call counters ``get_workspace`` creates; registered in its cache.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + + wss = [] + for maxt in ar_module.BUCKETS: + words = 4 * maxt * WORLD * (HIDDEN // 2) + handle = _make_mnnvl_mcast_buffer(dist.group.WORLD, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=dist.group.WORLD, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + wss.append(ws) + buckets = ar_module.McBuckets(wss) + ar_module._WORKSPACE[HIDDEN] = buckets + assert ar_module.get_workspace(mapping, HIDDEN) is buckets + return buckets + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + # The covered rows are 1..DRAFT_MC_AR_MAX_ROWS, every one of them driven; the row kernel's drafter form, the one + # the bitwise claim is about (the cluster form is exact-class). + covered = set(modeling.DRAFT_MC_AR.keys) + assert covered == set(range(1, modeling.DRAFT_MC_AR_MAX_ROWS + 1)), sorted(covered) + assert covered <= set(TOKENS), (sorted(covered), TOKENS) + assert any(m not in covered for m in TOKENS), ( + "drive at least one row count past the multicast boundary" + ) + assert not modeling.DRAFT_MC_AR_CLUSTER, ( + "the cluster drafter form is exact-class, not bitwise stock" + ) + # A declined multicast boundary logs (``missed``) and falls back instead of raising. + modeling.set_require(False) + _register_workspace(mapping) + mc = _McCalls(ar_module.m3_mc_allreduce_norm) + ar_module.m3_mc_allreduce_norm = mc + if RANK == 0: + print( + f"{RESULT_TAG} multicast drafter form at rows {sorted(covered)} (DRAFT_MC_AR_MAX_ROWS" + f" {modeling.DRAFT_MC_AR_MAX_ROWS}); built-in path at the others and when declined", + flush=True, + ) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + # The multicast drafter kernel takes no strategy: it sums the ranks in FP32 in rank order, which is stock's ONESHOT + # (the strategy a served drafter resolves AUTO to). Under NCCL its rows are therefore compared with stock ONESHOT; + # NCCL itself (which can round cancelling partials differently) stays the reference of the built-in rows. + oneshot = AllReduce(mapping=mapping, strategy=AllReduceStrategy.ONESHOT, dtype=torch.bfloat16) + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + for declined in (False, True): + path = "declined" if declined else "boundary" + ulp, res_ulp = 0.0, 0.0 + for i, m in enumerate(TOKENS): + xs, residual, w = build(m) + target_out, target_res, ran_mc = _target( + mapping, ar, xs[RANK], residual, w, NORM_KEYS[i % 2], mc, declined + ) + expect_mc = not declined and m in covered + stock_ar = oneshot if expect_mc and strategy == AllReduceStrategy.NCCL else ar + stock_out, stock_res = _stock(stock_ar, xs[RANK], residual, w) + torch.cuda.synchronize() + assert ran_mc == expect_mc, ( + f"{strategy.name} {path} {name} m={m}: the multicast drafter form " + f"{'ran' if ran_mc else 'did not run'}" + ) + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {path} ({'multicast' if ran_mc else 'built-in'}) {name}" + f" m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, path, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {path:8s} | {name:34s} | target = stock vs fp64:" + f" normed {ulp:.3f} | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | both | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD, modeling, ar_module + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the drafter runs at tp4, got {WORLD} ranks" + try: + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.modeling as _modeling # noqa: E501 + + modeling, ar_module = _modeling, _ar + assert modeling.m3_mc_allreduce is ar_module, ( + "the method calls the kernels module this patches" + ) + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..0b372cc02921 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,665 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``, 4 ranks. + +Its own launcher, as the catalog's collective matrices are: run plainly it spawns +one rank process per device in CUDA_VISIBLE_DEVICES and kills them all on the +first failure or on the deadline. A rank whose peers never push polls forever, +so the deadline is what keeps a broken exchange from wedging the calling run. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_mc_allreduce_norm_op_matrix.py + +Each rank builds the three multicast workspaces (8 / 16 / 48 tokens) the target +builds, over the torch.distributed group (the non-MPI path of +``_make_mnnvl_mcast_buffer``). Every rank can regenerate every rank's inputs from +the shared seed convention, which is what keeps each reference arithmetic rather +than a second collective. Producer pushes (``prepushed`` 2 / 3) are emulated by +storing the canonicalized rows through the multicast mapping at the call parity +the flags select, which is exactly what the producing kernels store. + +Not a pytest module, despite the ``check_*`` bodies: they are one fixed sequence +inside a single 4-rank job whose workspace state (call parities) carries from +one to the next. The collected entry point is +``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +finalize_entry: Any = None +ar_module: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +# ---------------------------------------------------------------- inputs + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + """One rank's [m, H] BF16 rows; ``rank`` -1 for tensors every rank holds alike.""" + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return _rows(seed + 7, -1, m, 4.0) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _finalize_inputs(seed: int, rank: int, m: int, topk: int, dropped: int): + """One rank's unfinalized MoE outputs: fc2 [m * topk, H], perm [m, topk], fp32 scales, shared [m, H].""" + g = _gen(seed + 23, rank) + rows = m * topk + fc2 = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + perm = torch.randperm(rows, generator=g, device="cuda").to(torch.int32) + if dropped: + drop = torch.randperm(rows, generator=g, device="cuda")[:dropped] + perm[drop] = -1 + perm = perm.view(m, topk).contiguous() + scales = torch.rand(m, topk, generator=g, device="cuda", dtype=torch.float32) + shared = _rows(seed + 29, rank, m, 0.5) + return fc2, perm, scales, shared + + +# ---------------------------------------------------------------- workspace / producer emulation + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor, exchange_set: int) -> None: + """Store this rank's rows into every rank's exchange, as a producer kernel does before a prepushed call. + + Slot of (token t, cluster CTA c): set * 2 * half + parity(t, c) * half + (t * world + rank) * H / 2 + c * H / 16 + words, parity read from the CTA's flag; -0.0 halves are stored as +0.0 (0x80000000 is the empty word). + """ + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = ( + exchange_set * 2 * half + + parity.view(m, CLUSTER, 1) * half + + (t * WORLD + RANK) * (HIDDEN // 2) + + c * chunk + + w + ) + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +# ---------------------------------------------------------------- gates + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +#: Share of normed elements one BF16 step off the float64 reference: measured at most 1.1e-4 per call over +#: every call of this matrix on GB300 (see ``_assert_normed``). The gate sits ~18x above that and far below what a +#: wrong statistic produces (most elements move). +NORMED_STEP_FRACTION = 0.002 +NORMED_STEP_SEEN: List[float] = [] + + +def _assert_normed(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """At most one BF16 rounding step from the float64 reference, and rarely that. + + The kernel's RMS statistic is a shuffle / warp / cluster tree of FP32 partial sums and its ``rsqrtf`` is + approximate (~2^-22 relative): the scaled value moves by ~1e-6 relative, which changes the BF16 rounding only + for values within that distance of a rounding midpoint -- one step, never two. The fraction bound catches a + statistic that is systematically off (a missing rank, a wrong H), which moves most elements. + """ + assert out.dtype == ref.dtype == torch.bfloat16 and out.shape == ref.shape, where + o, r = out.float(), ref.float() + _, exponent = torch.frexp(torch.maximum(o.abs(), r.abs())) + step = torch.ldexp(torch.ones_like(o), exponent - 8) + diff = (o - r).abs() + assert bool((diff <= step).all()), f"{where}: max diff {diff.max().item()} beyond one BF16 step" + frac = (diff != 0).float().mean().item() + NORMED_STEP_SEEN.append(frac) + assert frac <= NORMED_STEP_FRACTION, f"{where}: {frac:.4f} of normed elements off by one step" + + +def _assert_quant(outs, ref, m: int, where: str) -> None: + """The MXFP8 copy: bit-exact against the recipe applied to the op's own normed rows (a pure function of + them), and equal to the reference's wherever the reference's normed 32-column block is the op's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + mxfp8_rows, + sf_written, + ) + + q, sf = outs[2], outs[3] + own_q, own_sf = mxfp8_rows(outs[0]) + written = sf_written(m, HIDDEN // 32, sf.device) + _assert_bitwise(q, own_q, f"{where} e4m3 rows") + _assert_bitwise(sf[written], own_sf[written], f"{where} UE8M0 scales") + same_block = (outs[0].view(m, -1, 32) == ref[0].view(m, -1, 32)).all(-1) + qb = q.view(torch.uint8).view(m, -1, 32) + rqb = ref[2].view(torch.uint8).view(m, -1, 32) + assert bool((qb == rqb).all(-1)[same_block].all()), ( + f"{where}: e4m3 differs on a block whose normed rows agree" + ) + + +def _check_outputs(outs, ref, m: int, quant: bool, where: str) -> None: + assert len(outs) == (4 if quant else 2), (where, len(outs)) + _assert_bitwise(outs[1], ref[1], f"{where} updated residual") + _assert_normed(outs[0], ref[0], f"{where} normed") + if quant: + _assert_quant(outs, ref, m, where) + + +# ---------------------------------------------------------------- calls + + +def _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, opts): + return allreduce_entry.m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + ws.max_tokens, + quant, + prepushed, + opts, + ) + + +def _plain_case(spec: Dict[str, Any], seed: int, where: str) -> None: + m, maxt, quant, prepushed = spec["m"], spec["maxt"], spec["quant"], spec["prepushed"] + ws = WORKSPACES[maxt] + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] if spec["x2"] else None + residual, norm_w = _residual(seed, m), _norm_w(seed) + x, x2 = xs[RANK], (None if x2s is None else x2s[RANK]) + if prepushed >= 2: + _push(ws, x, 0) + if prepushed == 3: + _push(ws, x2, 1) + with validating(allreduce_entry.m3_mc_allreduce_norm): + outs = _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, spec["opts"]) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, residual, norm_w, EPS, quant, spec.get("drafter", False) + ) + _check_outputs(outs, ref, m, quant, where) + dist.barrier() + + +# ---------------------------------------------------------------- checks + + +def check_plain_cells() -> None: + """Every cell of ``m3_mc_allreduce_norm.CELLS`` through ``__call__`` under ``validating``, gated on its + ``reference``.""" + for i, cell in enumerate(allreduce_entry.m3_mc_allreduce_norm.CELLS): + _plain_case(cell.spec, 1000 + 17 * i, cell.why) + + +def check_finalize_cells() -> None: + """Every cell of ``m3_mc_finalize_allreduce_norm.CELLS``, gated on its ``reference``.""" + op = finalize_entry.m3_mc_finalize_allreduce_norm + for i, cell in enumerate(op.CELLS): + spec = cell.spec + m, maxt, quant, topk = spec["m"], spec["maxt"], spec["quant"], spec["topk"] + seed = 2000 + 17 * i + ws = WORKSPACES[maxt] + per_rank = [_finalize_inputs(seed, r, m, topk, spec["dropped"]) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + with validating(op): + outs = op( + fc2, + perm, + scales, + shared, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + maxt, + quant, + spec["opts"], + ) + torch.cuda.synchronize() + ref = op.reference( + [p[0] for p in per_rank], + [p[1] for p in per_rank], + [p[2] for p in per_rank], + [p[3] for p in per_rank], + residual, + norm_w, + EPS, + quant, + ) + _check_outputs(outs, ref, m, quant, cell.why) + dist.barrier() + + +def check_the_target_wrappers_select_bucket_and_options() -> None: + """The target's own wrappers (bucket by token count, option table) on top of the entries, every token count of + the 8-token bucket plus one of each larger bucket, and the finalize wrapper; the MXFP8 publication is the + prequant slot's business and is left off.""" + buckets = ar_module.McBuckets([WORKSPACES[b] for b in BUCKETS]) + for m in (*range(1, 9), 11, 40): + seed = 3000 + m + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + normed, updated = ar_module.m3_mc_allreduce_norm( + xs[RANK], residual, norm_w, EPS, buckets, x2=x2s[RANK] + ) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, False) + _check_outputs([normed, updated], ref, m, False, f"wrapper m={m}") + dist.barrier() + for m in (9, 33): + seed = 3100 + m + per_rank = [_finalize_inputs(seed, r, m, 4, 2) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + normed, updated = ar_module.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, buckets + ) + torch.cuda.synchronize() + ref = finalize_entry.m3_mc_finalize_allreduce_norm.reference( + *[[p[j] for p in per_rank] for j in range(4)], residual, norm_w, EPS, False + ) + _check_outputs([normed, updated], ref, m, False, f"finalize wrapper m={m}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """Two boundaries captured in one graph (8- and 16-token buckets), replayed with fresh inputs. + + The decode graph holds every boundary call; each replay rotates the same per-CTA parities the eager calls do, + so a replay that reused a stale half or skipped a re-arm reads the previous replay's rows and fails the + bitwise residual gate (every replay gets new inputs). + """ + shapes = ((3, 8, 15), (12, 16, 15)) + norm_w = _norm_w(4000) + bufs = [] + for m, maxt, opts in shapes: + bufs.append( + dict( + m=m, + maxt=maxt, + opts=opts, + x=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + x2=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + residual=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + ) + + def body(): + return [ + _call_plain( + WORKSPACES[b["maxt"]], b["x"], b["x2"], b["residual"], norm_w, False, 0, b["opts"] + ) + for b in bufs + ] + + def fill(seed: int): + refs = [] + for j, b in enumerate(bufs): + s = seed + 31 * j + xs = [_rows(s, r, b["m"]) for r in range(WORLD)] + x2s = [_rows(s + 3, r, b["m"]) for r in range(WORLD)] + b["x"].copy_(xs[RANK]) + b["x2"].copy_(x2s[RANK]) + b["residual"].copy_(_residual(s, b["m"])) + refs.append( + allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, b["residual"].clone(), norm_w, EPS, False + ) + ) + torch.cuda.synchronize() + dist.barrier() + return refs + + refs = fill(4100) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm = body() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(warm[j], refs[j], b["m"], False, f"graph warm-up m={b['m']}") + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = body() + for rep in range(4): + refs = fill(4200 + 100 * rep) + graph.replay() + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(captured[j], refs[j], b["m"], False, f"replay {rep} m={b['m']}") + dist.barrier() + del graph + dist.barrier() + + +def check_the_gate_sees_a_missing_rank() -> None: + """The residual gate is not blind to the reduction: dropping any one rank's partial from the reference + moves most elements of the updated residual (each rank's rows are distinct by construction).""" + m, seed = 4, 5000 + xs = [_rows(seed, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + full = allreduce_entry.m3_mc_allreduce_norm.reference(xs, None, residual, norm_w, EPS, False) + for r in range(WORLD): + partial = [x if i != r else torch.zeros_like(x) for i, x in enumerate(xs)] + wrong = allreduce_entry.m3_mc_allreduce_norm.reference( + partial, None, residual, norm_w, EPS, False + ) + frac = (wrong[1] != full[1]).float().mean().item() + assert frac > 0.9, f"dropping rank {r} changes only {frac:.3f} of the residual" + dist.barrier() + + +def _expect_rejected(fn, match: str, where: str) -> None: + try: + fn() + except AssertionError as exc: + assert match in str(exc), f"{where}: rejected for the wrong reason: {exc}" + else: + raise AssertionError(f"{where}: the guard admitted it") + + +def check_guards() -> None: + """``is_valid`` refuses what the op accepts and answers wrongly. Nothing here reaches the kernel: the guard + raises first, on every rank alike, so no rank is left polling.""" + m, seed = 4, 6000 + ws8, ws16 = WORKSPACES[8], WORKSPACES[16] + x, residual, norm_w = _rows(seed, RANK, m), _residual(seed, m), _norm_w(seed) + op = allreduce_entry.m3_mc_allreduce_norm + + def call(ws_uc, flags, norm): + with validating(op): + op( + x, + None, + residual, + norm, + EPS, + ws_uc, + ws8.mc.data_ptr(), + flags, + RANK, + WORLD, + CLUSTER, + False, + 8, + False, + 0, + 15, + ) + + strided = torch.empty(2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[::2] + strided.copy_(norm_w) + _expect_rejected( + lambda: call(ws8.uc, ws8.flags_for(CLUSTER), strided), "contiguous", "strided norm_w" + ) + _expect_rejected( + lambda: call(ws8.uc[: ws8.uc.numel() // 2], ws8.flags_for(CLUSTER), norm_w), + "words", + "short workspace", + ) + _expect_rejected( + lambda: call(ws8.uc, ws16.flags_for(CLUSTER), norm_w), "counters", "another bucket's flags" + ) + + fin = finalize_entry.m3_mc_finalize_allreduce_norm + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, 9, 4, 0) + bad = perm.clone() + bad[0, 0] = fc2.shape[0] + + def call_fin(): + with validating(fin): + fin( + fc2, + bad, + scales, + shared, + _residual(seed, 9), + norm_w, + EPS, + ws16.uc, + ws16.mc.data_ptr(), + ws16.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + 16, + False, + 15, + ) + + _expect_rejected(call_fin, "perm spans", "perm past the FC2 rows") + dist.barrier() + + +CHECKS = ( + check_plain_cells, + check_finalize_cells, + check_the_target_wrappers_select_bucket_and_options, + check_cuda_graph_replays, + check_the_gate_sees_a_missing_rank, + check_guards, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, allreduce_entry, finalize_entry, ar_module, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the entries are certified at 4 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as _finalize, + ) + + allreduce_entry, finalize_entry, ar_module = _plain, _finalize, _ar + # Spawned by file path, not collected by pytest: reach the sibling helper the plain way (confined to the + # rank processes). /comm -> tests/unittest/_torch/modeling_v2. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peers' polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print( + f"[rank {RANK}] {len(CHECKS)} checks passed; normed one-step fraction max " + f"{max(NORMED_STEP_SEEN):.5f} over {len(NORMED_STEP_SEEN)} calls", + flush=True, + ) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_rank_job.py new file mode 100644 index 000000000000..1115ff3b5976 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_rank_job.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run this target's collective rank job from pytest, and report what it did. + +The same contract as ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``: the +launcher beside this file owns the per-device rank processes and the deadline +that kills a wedged collective; this only selects the devices, starts it in a +fresh interpreter (by file path: this tree is not a package) and turns its exit +code into an assertion. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +WORLD_SIZE = 4 +"""tp4's world size -- the topology these entries are certified for.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the launcher's own deadline, so its message wins the race.""" + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + if len(devices) < WORLD_SIZE: + pytest.skip( + f"this entry is certified at world size {WORLD_SIZE}; only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def _load_launcher_constants(launcher: Path): + """Read the launcher's module-level budget without importing it by name (it imports only torch at + module scope).""" + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str) -> str: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices; returns its stdout.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + + # The ranks import tensorrt_llm absolutely. From a source checkout that is not installed, put the repo root + # (tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root) on their path; from a tests + # tree run against an installed package, they import the installed one. + repo_root = Path(__file__).resolve().parents[8] + if (repo_root / "tensorrt_llm").is_dir(): + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p + ) + else: + assert importlib.util.find_spec("tensorrt_llm") is not None, ( + f"no tensorrt_llm/ at the repo root {repo_root} and no installed tensorrt_llm" + ) + + constants = _load_launcher_constants(launcher) + timeout = constants.DEADLINE_S + _LAUNCHER_GRACE_S + + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) + return completed.stdout diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_allreduce.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_allreduce.py new file mode 100644 index 000000000000..27d45f08edc4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_allreduce.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the ``m3_allreduce`` certification matrix. + +The matrix is ``_m3_allreduce_op_matrix.py`` beside this file, its own 4-rank launcher; see ``_rank_job``. It drives +every cell of the entry against its reference, then requires the fused draft loop's boundaries built on it (plain +all-reduce, then the fused add + RMSNorm) to equal stock Eagle3's bit for bit under every strategy, and reports both +against an fp64 reference as a warning, one line per strategy and case. +""" + +import warnings + +import _rank_job +import pytest +import torch + +pytestmark = [ + pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _rank_job.WORLD_SIZE, + reason=f"certified at world size {_rank_job.WORLD_SIZE}", + ), +] + +_RESULT_TAG = "[ulp]" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_allreduce_op_matrix() -> None: + out = _rank_job.run("m3_allreduce") + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + "draft boundary (= stock Eagle3) vs fp64, max BF16 ulps\n" + "\n".join(lines), stacklevel=1 + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py new file mode 100644 index 000000000000..ad7539d27829 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for CAPTURE_NORM_IN_BOUNDARY on the tp4 multicast all-reduce + RMSNorm op. + +The matrix is ``_m3_capture_in_boundary_op_matrix.py`` beside this file, its own 4-rank launcher; see ``_rank_job``. +It requires the boundary's capture output to equal ``m3_spec_tail.capture_norm`` of its updated residual bit for bit, +on both kernels the tp4 op runs (the oneshot cluster kernel, the row kernel), and the op's other outputs to be +unchanged by it. +""" + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "the multicast all-reduce requires CUDA devices" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_capture_norm_in_boundary_is_capture_norm() -> None: + out = _rank_job.run("m3_capture_in_boundary") + assert "[case]" in out, f"the launcher ran no case:\n{out}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py new file mode 100644 index 000000000000..3cc1592bc3c5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the fast draft loop's layer boundaries against stock Eagle3's. + +The matrix is ``_m3_draft_ar_norm_op_matrix.py`` beside this file, its own 4-rank launcher; see ``_rank_job``. It +requires the target's drafter boundary (``modeling.M3Eagle3Worker._m3_draft_ar_norm``: the multicast all-reduce + +residual + RMSNorm in its drafter form at the rows ``DRAFT_MC_AR`` covers, else and when that declines +``M3Eagle3Draft.ar_norm`` on the strategy's all-reduce) to equal stock Eagle3's boundary (plain all-reduce, then the +``RMSNorm`` module with the residual) bit for bit under every strategy, and reports both against an fp64 reference as a +warning, one line per strategy, path and case. +""" + +import warnings + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_RESULT_TAG = "[ulp]" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_draft_ar_norm_is_stock_eagle3() -> None: + out = _rank_job.run("m3_draft_ar_norm") + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + "draft boundary (= stock Eagle3) vs fp64, max BF16 ulps\n" + "\n".join(lines), stacklevel=1 + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..31a054642e89 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the multicast all-reduce + RMSNorm entries' certification matrix. + +Covers both ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm`` (one +kernel, one exchange protocol). The matrix is ``_m3_mc_allreduce_norm_op_matrix.py`` +beside this file, its own 4-rank launcher; see ``_rank_job`` for why. +""" + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_mc_allreduce_norm_op_matrix() -> None: + _rank_job.run("m3_mc_allreduce_norm") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..7cf1ffb5a141 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests of the ``m3_mc_finalize_allreduce_norm`` entry that need no collective. + +Its cells are certified, with ``m3_mc_allreduce_norm``'s (one kernel, one exchange protocol), by the 4-rank matrix +``_m3_mc_allreduce_norm_op_matrix.py`` that ``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py`` runs. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + finalize_partial, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as op, +) + + +def test_finalize_partial_rounds_the_routed_sum_before_the_shared_expert() -> None: + """As stock: the finalized routed output is BF16, and the shared expert is added to it (shared.add_(routed)). + Four routed rows of 1.0 at FP32 routing weight 0.2509765625 sum to 1.00390625, which rounds to BF16 1.0; plus the + shared -1 that is 0, where adding before the rounding would leave 0.00390625.""" + fc2 = torch.ones((4, 8), dtype=torch.bfloat16) + perm = torch.arange(4, dtype=torch.int32).view(1, 4) + scales = torch.full((1, 4), 0.2509765625, dtype=torch.float32) + shared = torch.full((1, 8), -1.0, dtype=torch.bfloat16) + out = finalize_partial(fc2, perm, scales, shared) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(out, shared.clone().add_(routed)) + assert torch.equal(out, torch.zeros(1, 8, dtype=torch.bfloat16)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_each_cells_partial_is_the_stock_finalize_then_shared_add(cell) -> None: + """Per cell shape (rows, top-k, dropped slots), every rank's partial in the entry's reference is the stock MoE's: + the finalized routed rows rounded to BF16, then ``shared.add_(routed)``. The 4-rank matrix drives the cells on the + GPU against this reference.""" + spec = cell.spec + g = torch.Generator().manual_seed(spec["m"] * 31 + spec["dropped"]) + m, topk, hidden = spec["m"], spec["topk"], 64 + fc2 = torch.randn(m * topk, hidden, generator=g).to(torch.bfloat16) + perm = torch.randperm(m * topk, generator=g).view(m, topk).to(torch.int32) + perm.view(-1)[: spec["dropped"]] = -1 + scales = torch.rand(m, topk, generator=g) + shared = torch.randn(m, hidden, generator=g).to(torch.bfloat16) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(finalize_partial(fc2, perm, scales, shared), shared.clone().add_(routed)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..4258edbc4638 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_embedding entry of the MiniMax-M3 tp4 target: an exact gather with zero rows outside the +table.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.embedding.m3_embedding import ( # noqa: E501 + m3_embedding as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" + +VOCAB, HIDDEN = 4096, 6144 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + ids = torch.randint( + 0, VOCAB, (spec["rows"],), generator=g, device="cuda", dtype=spec["id_dtype"] + ) + if spec["pad"]: + ids[1::3] = -1 + with validating(op): + out = op(ids, table) + assert torch.equal(out, op.reference(ids, table)) + if spec["pad"]: + assert not out[1::3].any(), "padded ids must embed to zero rows" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py new file mode 100644 index 000000000000..e44ad56d743d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" + +import torch + + +def swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): + """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" + v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) + return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..f17d7ec63189 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_bf16_linear entry of the MiniMax-M3 tp4 target (the drafter's projections).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_bf16_linear import ( # noqa: E501 + m3_bf16_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the + row, RMS well under one; measured on GB300: worst 0.5 unit.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w = (torch.randn(spec["n"], spec["k"], generator=g, device="cuda") * 0.02).to(torch.bfloat16) + with validating(op): + out = op(x, w) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(x, w), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py new file mode 100644 index 000000000000..c53155b7c02b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise oracle of the dense MLP's swigluoai C++ op (ops/gemm/m3_dense_swiglu, +trtllm::minimax_m3_sm107_tp4_swiglu_oai) against the Triton silu_and_mul kernel the dense MLP runs today +(modules/swiglu.py at limit 7, alpha 1.702, beta 1): every bf16 gate bit pattern against special and random up values, +every up bit pattern against special and random gates, random pairs, and the dense MLP's shapes at every decode row +count, eager and replayed from a CUDA graph, with and without the programmatic launch.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.silu_and_mul +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_dense_swiglu import ( # noqa: E501 + ALPHA, + BETA, + LIMIT, + m3_dense_swiglu, +) + +assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp4_swiglu_oai"): + pytest.skip( + "needs the TP4 op library with minimax_m3_sm107_tp4_swiglu_oai", allow_module_level=True + ) + +# The TP4 dense MLP's activation width (gate_up 3072 -> d 1536). +DENSE_D = dispatch.DENSE_DOWN[1] + + +def _triton(x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=LIMIT, swiglu_alpha=ALPHA, swiglu_beta=BETA + ) + + +def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" + return out + + +def _row(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + return torch.cat([gate.view(1, -1), up.view(1, -1)], dim=1).contiguous() + + +def _all_bf16() -> torch.Tensor: + return ( + torch.arange(1 << 16, dtype=torch.int32, device="cuda").to(torch.int16).view(torch.bfloat16) + ) + + +def _specials(n: int, seed: int) -> list: + g = torch.Generator(device="cuda").manual_seed(seed) + vals = [0.0, -0.0, 1.0, -1.0, 6.99, 7.0, 7.01, -7.0, -7.01, 1e-38, -1e-38, 1e-40, 88.0, -88.0, 3e38, -3e38, + float("inf"), float("-inf"), float("nan")] # fmt: skip + rows = [torch.full((n,), v, dtype=torch.bfloat16, device="cuda") for v in vals] + for scale in (0.05, 1.0, 4.0, 30.0): + rows.append((scale * torch.randn(n, generator=g, device="cuda")).to(torch.bfloat16)) + rows.append( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + return rows + + +def _assert_bitwise(got: torch.Tensor, ref: torch.Tensor, what: str) -> None: + bad = (got.view(torch.int16) != ref.view(torch.int16)).nonzero().flatten() + assert bad.numel() == 0, f"{what}: {bad.numel()} mismatches, first at {bad[:4].tolist()}" + + +def test_every_gate_pattern_is_bitwise_the_triton_kernel() -> None: + gate = _all_bf16() + for j, up in enumerate(_specials(gate.numel(), 1)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"up set {j}") + + +def test_every_up_pattern_is_bitwise_the_triton_kernel() -> None: + up = _all_bf16() + for j, gate in enumerate(_specials(up.numel(), 2)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"gate set {j}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: + g = torch.Generator(device="cuda").manual_seed(100 + seed) + n = 1 << 22 + + def pattern() -> torch.Tensor: + bits = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + return bits.to(torch.int16).view(torch.bfloat16) + + x = _row(pattern(), pattern()) + _assert_bitwise(_op(x), _triton(x), f"seed {seed}") + + +@pytest.mark.parametrize("pdl", [True, False], ids=["pdl", "plain"]) +@pytest.mark.parametrize("rows", range(1, dispatch.DECODE_MAX_TOKENS + 1)) +def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pdl: bool) -> None: + """The TP4 dense MLP's gate_up rows (a GEMV output's range plus values past the clamps), eager and replayed from a + CUDA graph after fresh rows are written in place, against the Triton kernel on the same rows.""" + g = torch.Generator(device="cuda").manual_seed(2000 + rows) + x = (3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, DENSE_D : DENSE_D + 8] = x[:, :8] + _assert_bitwise(_op(x, pdl), _triton(x), f"eager rows {rows}") + out = torch.empty(rows, DENSE_D, dtype=torch.bfloat16, device="cuda") + _op(x, pdl) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out.copy_(_op(x, pdl)) + x.copy_((3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(out, _triton(x), f"graph rows {rows}") + + +def test_the_op_declines_other_parameters_and_layouts() -> None: + x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + + +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: + """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton + kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" + rows, d = cell.spec["rows"], cell.spec["d"] + g = torch.Generator(device="cuda").manual_seed(3000 + rows) + x = (3.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, d : d + 8] = x[:, :8] + _assert_bitwise(_op(x), _triton(x), cell.why) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py new file mode 100644 index 000000000000..ddf72d7556e9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_draft_lowm entry: the Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M +kernels (direct SIMT and split-K tensor-core), at the TP4 drafter shards, at the rows its table certifies. + +Every cell is held to the fp64 reference within the derived band. The drafter's projections (``M3Eagle3Draft._linear`` +over ``kernels/m3_draft_lowm``) are held on the drafter's weight table: a covered row count is the entry's call bit for +bit, an uncovered one the plain GEMM's. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_draft_lowm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + TACTICS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 + m3_draft_lowm as op, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" + +# Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). +# +# Kernel and reference multiply the same BF16 operands (every product exact in fp32); the kernel sums in fp32 (the +# direct kernel one dot product in K order, the split-K kernel tensor-core blocks and then its split partials), the +# reference in fp64, and both round once to BF16. With w ~ N(0, 1 / K) and x ~ N(0, 1) the outputs are O(1) and the +# fp32 order difference is ~sqrt(K) * 2^-24 of the summed magnitudes, a few 1e-6 at K = 18432: far below half a BF16 +# ulp, so the two roundings disagree by at most one ulp of the element -- up to two units of the row scale for an +# element in its row maximum's binade -- and only where the sum sits that close to a rounding boundary, a few in ten +# thousand elements, hence 0.1 ulp RMS. The same derivation as the decode GEMV's band (test_modeling_v2_m3_mx_gemv), +# with no slack under it: a second ulp would mean a wrong sum. +ELEMENT_ULP = 2.0 +RMS_ULP = 0.1 + +# (N, K) -> the drafter's weight key (M3Eagle3Draft.w). +DRAFTER_LINEARS = { + (2304, 12288): "qkv", + (6144, 2048): "o", + (9216, 6144): "gu", + (6144, 4608): "dn", + (6144, 18432): "fc", +} + + +@pytest.fixture(autouse=True) +def _policy(): + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + return x, w + + +def test_the_table_is_the_tp4_drafter() -> None: + assert set(TACTICS) == set(DRAFTER_LINEARS) + assert {c.spec["nk"] for c in op.CELLS} == set(TACTICS) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + kind, tactic = tactic_for(n, k, m) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + args = ( + x, + w, + out, + kind, + tactic, + True, + ) # the target launches with PDL when the environment enables it + certifying(op) + assert op(*args) is out + ref = op.reference(*args) + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=RMS_ULP) + # PDL only changes the launch attribute. + plain = torch.empty_like(out) + op(x, w, plain, kind, tactic, False) + assert torch.equal(plain.view(torch.int16), out.view(torch.int16)) + + +def _draft(linears=DRAFTER_LINEARS): + """A drafter holding only its projection weights (the module is not built).""" + draft = object.__new__(dispatch.M3Eagle3Draft) + draft.w = {key: build(*nk, 1, seed=nk[0] + nk[1])[1] for nk, key in linears.items()} + draft._lowm_keys = frozenset() + return draft + + +@pytest.fixture(scope="module") +def prebuilt(): + draft = _draft() + draft.prebuild_linears() # 0 when an earlier test in this process compiled them + assert draft._lowm_keys == frozenset(DRAFTER_LINEARS.values()) + assert draft.prebuild_linears() == 0, "a second prebuild finds nothing left to compile" + return draft + + +@pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) +def test_the_projection_is_the_entry_and_falls_back_past_its_rows(prebuilt, nk) -> None: + n, k = nk + key = DRAFTER_LINEARS[nk] + w = prebuilt.w[key] + use_pdl = get_env_enable_pdl() + dispatch.set_require(True) + for m in ROWS[nk]: + assert dispatch.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt._linear(x, key) + kind, tactic = tactic_for(n, k, m) + want = op(x, w, torch.empty(m, n, dtype=x.dtype, device="cuda"), kind, tactic, use_pdl) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)), m + # Rows the table does not certify (past its rows): the plain GEMM, silently, even under require. + for m in sorted(set(range(1, dispatch.DECODE_MAX_TOKENS + 2)) - set(ROWS[nk])): + assert not dispatch.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt._linear(x, key) + assert torch.equal(got.view(torch.int16), F.linear(x, w).view(torch.int16)), m + + +def test_a_projection_outside_the_table_is_the_plain_gemm() -> None: + assert not m3_draft_lowm.installed(torch.empty(256, 512, dtype=torch.bfloat16)) + draft = _draft({(256, 512): "o"}) + draft._lowm_keys = frozenset(k for k in ("o",) if m3_draft_lowm.installed(draft.w[k])) + x, _ = build(256, 512, 1, seed=3) + got = draft._linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_a_covered_call_not_prebuilt_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + monkeypatch.setattr(m3_draft_lowm, "_KEYS", frozenset(dispatch.DRAFT_LOWM.keys)) + w = build(6144, 2048, 1, seed=4)[1] + x, _ = build(6144, 2048, 1, seed=5) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not compiled before warmup"): + m3_draft_lowm.draft_linear(w, x, None) + # Outside require the drafter falls back to the plain GEMM. + dispatch.set_require(False) + assert m3_draft_lowm.draft_linear(w, x, None) is None + + +def test_guard_refuses_a_strided_activation() -> None: + x, w = build(6144, 2048, 2, seed=6) + wide = torch.zeros(2, 4096, dtype=torch.bfloat16, device="cuda") + wide[:, :2048] = x + kind, tactic = tactic_for(6144, 2048, 2) + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op( + wide[:, :2048], + w, + torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda"), + kind, + tactic, + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..fc71fc9f522a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp32_linear entry of the MiniMax-M3 tp4 target (the router beyond the decode routers' rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """torch's default fp32 matmul precision on this hardware may be TF32 (~2^-11 relative per product): the band is + stated in fp32 ulps of each row's largest logit and covers it (K = 6144 TF32 products: ~2^13 ulps worst element); + measured on GB300 with the engine's defaults: worst element ~1700 ulps, RMS ~300.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(spec["e"], spec["k"], generator=g, device="cuda") * 0.02 + with validating(op): + out = op(x, w) + assert out.dtype == torch.float32 and out.shape == (spec["m"], spec["e"]) + assert_within_ulp(out, op.reference(x, w), element_ulp=8192.0, rms_ulp=2048.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..03197b28202b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv entry: the decode MXFP8 GEMV on a producer-quantized activation. + +Operands as the target holds them: the weight through mxfp8_quantize with the swizzled scale layout (the MXFP8 linear +method's storage), the activation through mxfp8_quantize with the swizzled layout (what the boundary all-reduce and +the attention merge publish). +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" + + +# Element band, in ULP units of the row scale: see compare. +ELEMENT_ULP = 1.5 + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Element ELEMENT_ULP of the row scale, relative RMS 0.1 ulp. + + Kernel and reference multiply bit-identical dequantized operands; the kernel sums in fp32 (tensor-core blocks, + then the split-K partials), the reference in fp64. Over K <= 6144 that order difference is ~1e-6 relative, so the + two bf16 roundings disagree only where the sum sits within that of a rounding boundary: one bf16 ulp of that + element, which for an element in its row maximum's binade is up to two of ``assert_within_ulp``'s units of the row + scale, and a few in ten thousand elements in aggregate. Any fp32 sum shows it: torch's IEEE fp32 matmul of the + same dequantized operands measured at most 1.22 units against the fp64 reference over every row count 1..48 on + the o_proj and the QKV (6 seeds each, GB300), the kernel at most 1.25, both at most 0.009 ulp RMS. The band + leaves the margin between that and a full extra ulp. + """ + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=0.1) + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + return x, xq.view(m, k), xsf, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + _, xq, xsf, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(xq, xsf, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + compare(out, op.reference(xq, xsf, w, wsf)) + + +def test_an_unserved_shape_returns_none() -> None: + """Shapes outside the tuned table fall back at the call site; the op says so rather than computing.""" + _, xq, xsf, w, wsf = build(256, 512, 1, seed=1) + assert op(xq, xsf, w, wsf) is None + + +def test_guard_refuses_linear_activation_scales() -> None: + """A linear [M, K / 32] scale buffer is read as the swizzled one.""" + x, _, _, w, wsf = build(6144, 2048, 4, seed=2) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, False) + with pytest.raises(AssertionError, match="swizzled"): + with validating(op): + op(xq.view(4, 2048), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + _, xq, xsf, w, wsf = build(6144, 2048, 1, seed=3) + wide = torch.zeros(6144, 4096, dtype=w.dtype, device="cuda") + wide[:, :2048] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :2048], wsf) + + +@pytest.mark.parametrize("nk", [(2560, 6144), (2304, 6144)], ids=["qkv_sparse", "qkv_dense"]) +@pytest.mark.parametrize("m", range(1, 9)) +def test_qkv_split8_config(nk, m) -> None: + """The TP4 dispatch's (split, ring) override for the QKV at decode rows: within the band of the fp64 reference, + and the graph replay bit-identical to the eager call (one split-K order per config).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as dispatch, + ) + + cfg = dispatch.mx_gemv_cfg(nk[0], nk[1], m) + assert cfg == (8, 8) + _, xq, xsf, wq, wsf = build(nk[0], nk[1], m, zlib.crc32(f"qkv8-{nk}-{m}".encode())) + eager = op(xq, xsf, wq, wsf, cfg=cfg) + assert eager is not None, "the GEMV declined the split-8 config" + compare(eager, op.reference(xq, xsf, wq, wsf)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(xq, xsf, wq, wsf, cfg=cfg) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured.view(torch.int16), eager.view(torch.int16)) + + +def test_qkv_split8_config_is_bounded(monkeypatch) -> None: + """The overrides cover only the two QKV projections (split 8) and the o_proj (strided k-tiles, rule 1b) at decode + rows; every other call, and every call with its kill knob off, keeps the shared table.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as dispatch, + ) + + assert dispatch.mx_gemv_cfg(2560, 6144, 9) is None + assert dispatch.mx_gemv_cfg(6144, 2048, 4) == (2, 8, 0) + assert dispatch.mx_gemv_cfg(6144, 2048, 9) is None + assert dispatch.mx_gemv_cfg(6144, 3072, 4) is None + monkeypatch.setattr(dispatch, "GEMV_OPROJ_STRIDED_TP4", False) + assert dispatch.mx_gemv_cfg(6144, 2048, 4) is None + assert dispatch.mx_gemv_cfg(2560, 6144, 4) == (8, 8) + + +@pytest.mark.parametrize( + "nk", + [(2560, 6144), (2304, 6144), (6144, 2048)], + ids=["qkv_sparse", "qkv_dense", "o_proj_strided"], +) +def test_prebuild_covers_the_split8_config_while_frozen(nk) -> None: + """The core's prebuild (``prebuild_kernels``: ``_gemv_variants`` + ``_precompile_cfg``) compiles every (rows, cfg) + ``_linear`` can pick for this projection, the TP4 (split, ring) override included: with compiles frozen afterwards + (as the target freezes them before warmup), every decode-row MXFP8 call at the override is served, none + declines.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as dispatch, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.mx_gemv import ( # noqa: E501 + op as mx_gemv, + ) + + n, k = nk + _, _, _, wq, wsf = build(n, k, 1, zlib.crc32(f"prebuild-{nk}".encode())) + mx_gemv.freeze(False) + mx_gemv.ensure_compiled(wq, wsf, True) + for rows, x_bf16, ar in dispatch._gemv_variants(n, k, None, 4): + cfg = None if x_bf16 else dispatch.mx_gemv_cfg(n, k, rows) + if cfg is not None: + dispatch._precompile_cfg(wq, wsf, rows, cfg, True, ar=ar) + else: + mx_gemv.precompile(wq, wsf, rows, x_bf16, use_pdl=True, ar=ar) + picked = [ + (rows, dispatch.mx_gemv_cfg(n, k, rows)) + for rows in sorted(dispatch.MX_GEMV_ROWS[nk]) + if dispatch.mx_gemv_cfg(n, k, rows) is not None + ] + assert picked, "the override covers decode rows of this projection" + for rows, cfg in picked: + resolved = mx_gemv._resolve(rows, n, k, False, True, cfg, True, None) + assert resolved is not None and mx_gemv.is_compiled(resolved[0]), (rows, cfg) + mx_gemv.freeze(True) + try: + for rows, cfg in picked: + _, xq, xsf, _, _ = build(n, k, rows, rows) + out = mx_gemv.m3_mx_gemv(xq, xsf, wq, wsf, cfg=cfg, use_pdl=True) + assert out is not None, (rows, cfg, mx_gemv.last_uncompiled) + finally: + mx_gemv.freeze(False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..2c3947b5cc81 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv_bf16 entry: the decode MXFP8 GEMV quantizing a bf16 activation in its prologue. + +The weight as the MXFP8 linear method stores it (mxfp8_quantize, swizzled scales). Besides the reference band, the +prologue's quantization is held to mxfp8_quantize itself: the GEMV on a bf16 x must equal, bit for bit, the +prequantized GEMV on mxfp8_quantize(x) -- the claim the target relies on when it mixes the two paths. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv_bf16 import ( # noqa: E501 + m3_mx_gemv_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" + +# The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. +ELEMENT_ULP = 1.5 + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return x, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(x, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + # The band of test_modeling_v2_m3_mx_gemv.compare: the quantized operands are the reference's exactly (checked + # below), the sums differ only in order. + assert_within_ulp(out, op.reference(x, w, wsf), element_ulp=ELEMENT_ULP, rms_ulp=0.1) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + pre = m3_mx_gemv(xq.view(m, k), xsf, w, wsf, use_pdl=True) + assert pre is not None and torch.equal(out.view(torch.int16), pre.view(torch.int16)), ( + "the prologue's MXFP8 quantization differs from mxfp8_quantize" + ) + + +def test_a_strided_activation_is_declined() -> None: + """The prologue reads x densely; the op declines a strided x (None) rather than computing on it.""" + x, w, wsf = build(6144, 2048, 2, seed=1) + wide = torch.zeros(2, 4096, dtype=torch.bfloat16, device="cuda") + wide[:, :2048] = x + assert op(wide[:, :2048], w, wsf) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_order.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_order.py new file mode 100644 index 000000000000..51cd471a9089 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_order.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The TP4 decode GEMV order overrides (modeling.mx_gemv_cfg: QKV split 8, o_proj strided k-tiles): +for every overridden (projection, rows) the output is the same on 3 eager calls and on a CUDA graph replay +(deterministic), and its max abs / max rel error vs the fp64 product of the dequantized operands stays within the +shipped order's band (the shared table's config on the same operands, over several seeds).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, + QKV, + QKV_SPARSE, + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 + mxfp8_dequant, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +CASES = [(nk, m) for nk in (QKV_SPARSE, QKV, O_PROJ) for m in (1, 2, 4, 8)] + + +def _err(y, ref): + d = (y.double() - ref).abs() + return float(d.max()), float((d / ref.abs().amax(dim=1, keepdim=True).clamp_min(1e-30)).max()) + + +@pytest.mark.parametrize("nk, m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +def test_order_override_is_deterministic_and_in_band(nk, m) -> None: + n, k = nk + cfg = dispatch.mx_gemv_cfg(n, k, m) + assert cfg is not None, "the override is off or does not cover this call" + worst_new, worst_ship = [0.0, 0.0], [0.0, 0.0] + for seed in range(4): + g = torch.Generator(device="cuda").manual_seed(1000 * seed + m) + w = (torch.randn(n, k, device="cuda", generator=g) * 2).to(torch.float8_e4m3fn) + ws = torch.randint(118, 124, (n * k // 32,), dtype=torch.uint8, device="cuda", generator=g) + x = (torch.randn(m, k, device="cuda", generator=g) * (0.5 + seed)).to(torch.bfloat16) + q, s = torch.ops.trtllm.mxfp8_quantize(x, True) + ref = ( + mxfp8_dequant(q.view(torch.uint8), s).double() + @ mxfp8_dequant(w.view(torch.uint8), ws).double().t() + ) + ys = [m3_mx_gemv(q, s, w, ws, True, cfg, True, None) for _ in range(3)] + assert all(y is not None for y in ys), "the GEMV declined the override" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + out = [] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out.append(m3_mx_gemv(q, s, w, ws, True, cfg, True, None)) + graph.replay() + torch.cuda.synchronize() + for y in ys[1:] + [out[0]]: + assert torch.equal(ys[0].view(torch.int16), y.view(torch.int16)), ( + "not deterministic (eager / graph)" + ) + ship = m3_mx_gemv(q, s, w, ws, True, None, True, None) + torch.cuda.synchronize() + for acc, y in ((worst_new, ys[0]), (worst_ship, ship)): + e = _err(y, ref) + acc[0], acc[1] = max(acc[0], e[0]), max(acc[1], e[1]) + print( + f"\nR1B1 tp4 {n}x{k} m={m} cfg {cfg}: max abs / rel vs fp64 {worst_new[0]:.3e} / {worst_new[1]:.3e} | shipped " + f"order {worst_ship[0]:.3e} / {worst_ship[1]:.3e}", + flush=True, + ) + # The band: the shipped order's own error (bf16 output rounding dominates), with 25 % headroom. + assert worst_new[0] <= 1.25 * worst_ship[0] and worst_new[1] <= 1.25 * worst_ship[1] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..54379ac1f3e2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_gemm entry of the MiniMax-M3 tp4 target (the decode-graph GEMM on a prequantized +activation).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most + 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, False), + element_ulp=2.0, + rms_ulp=0.25, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..b869ca3aef13 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_linear entry of the MiniMax-M3 tp4 target: MXFP8 quantize + GEMM on each backend the +engine selects (native CUTLASS, FlashInfer, the per-bucket tuned quantize + GEMM).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" + +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and + one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" + spec = cell.spec + if ( + spec["path"] == "flashinfer" + and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None + ): + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + args = (x, w, w_sf, alpha, torch.bfloat16, *_FLAGS[spec["path"]]) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py new file mode 100644 index 000000000000..a2ae95a27105 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_pinned_gemm entry (MX-TAC): the pinned decode-graph MXFP8 GEMM gives the same bits as +every candidate stock's AutoTuner can pick (CuTe DSL, CUTLASS), the same on every call, within the MXFP8 band of the +dequantized product.""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_pinned_gemm import ( # noqa: E501 + m3_mxfp8_pinned_gemm as op, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_pinned_gemm requires a CUDA device" + +TACTIC = modeling.MXFP8_DECODE_GEMM_TACTIC + +# The TP4 projections at rows the GEMVs do not serve (SHARED_DOWN 17..32, over 48) and the bf16 activations over the +# GEMVs' bf16 bounds. +CASES = [ + ((6144, 768), 17), + ((6144, 768), 32), + ((2560, 6144), 24), + ((2304, 6144), 64), + ((6144, 2048), 49), + ((1536, 6144), 20), + ((6144, 6144), 12), + ((6144, 3072), 56), +] + + +def test_the_pinned_tactic_is_cute_dsl() -> None: + assert TACTIC == FlashInferMXFP8GemmRunner.CUTE_DSL + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, TACTIC) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, TACTIC), + element_ulp=2.0, + rms_ulp=0.25, + ) + + +@pytest.mark.parametrize("nk,m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_pinned_decision_matches_every_stock_candidate(nk, m, scale) -> None: + n, k = nk + g = torch.Generator(device="cuda").manual_seed(n + k + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = (torch.randn(m, k, generator=g, device="cuda") * scale).to(torch.bfloat16) + x[0, ::32] *= 1e3 # one outlier per 32-element block on the first row + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + wq = wq.view(n, k) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + xq = xq.view(m, k) + out = op(xq, xsf, wq, wsf, torch.bfloat16, TACTIC) + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + for tactic in (FlashInferMXFP8GemmRunner.CUTE_DSL, FlashInferMXFP8GemmRunner.CUTLASS): + ref = runner.forward([xq, xsf, wq, wsf], tactic=tactic) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), tactic + again = op(xq, xsf, wq, wsf, torch.bfloat16, TACTIC) + assert torch.equal(out.view(torch.int16), again.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py new file mode 100644 index 000000000000..22e987e17979 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_pinned_quantize entry (MX-TAC): the pinned decode-graph MXFP8 activation quantize is the +stock quantize bit for bit, at every tactic stock's AutoTuner can pick.""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_pinned_quantize import ( # noqa: E501 + m3_mxfp8_pinned_quantize as op, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8QuantizeRunner + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_pinned_quantize requires a CUDA device" + +TACTIC = modeling.MXFP8_DECODE_QUANT_TACTIC + + +def test_the_pinned_tactic_is_cute_dsl() -> None: + assert TACTIC == MXFP8QuantizeRunner.CUTE_DSL + + +def _check(x: torch.Tensor) -> None: + m, k = x.shape + with validating(op): + xq, xsf = op(x, TACTIC) + tq, tsf = op.reference(x, TACTIC) + assert torch.equal(xq.view(torch.uint8).view(m, k), tq.view(torch.uint8).view(m, k)) + assert torch.equal(xsf.view(torch.uint8).flatten(), tsf.view(torch.uint8).flatten()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_certified_cells(cell, scale) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * scale).to(torch.bfloat16) + x[0, ::32] *= 1e3 # one outlier per 32-element block on the first row + _check(x) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py new file mode 100644 index 000000000000..c1c7aad25bce --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_target_head entry: the target lm_head shard GEMV over the lossless 12-bit coding of its +BF16 weight, against its reference, the in-tree BF16 head GEMM (``F.linear``) and fp64. + +* The coding round-trips bit for bit (code values and exception lists), and a corrupted coding fails the check. +* Integer-valued data (fp32 sums exact in any order): the kernel's logits equal the in-tree GEMM's bit for bit. +* Random data: the kernel's error against fp64 is the in-tree GEMM's (the same products, another fp32 order, one + BF16 rounding), its argmax is the in-tree one and its result is the same run to run. +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.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_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + CTA_WARPS, + HIDDEN, + TARGET_SHARD, + decode, + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + m3_target_head as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_target_head requires a CUDA device" + + +def _seed(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def _num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def _weight(n: int, k: int, seed: int, exceptions: bool = False) -> torch.Tensor: + """A [n, k] head of trained-weight scale; with ``exceptions``, values on both sides of the exponent window, exact + zeros and a large outlier in many tiles.""" + g = torch.Generator(device="cuda").manual_seed(seed) + w = torch.randn(n, k, device="cuda", generator=g) * 0.02 + if exceptions: + w[::97, ::131] = 1e-9 + w[::89, 7] = 0.0 + w[3, 5] = 7.5 + w[n - 1, k - 1] = -3.0 + return w.to(torch.bfloat16) + + +_CODED: dict = {} + + +def _coded(n: int, k: int, seed: int, exceptions: bool): + """(weight, encode_chunked(weight)), built once per shape.""" + key = (n, k, seed, exceptions) + if key not in _CODED: + _CODED.clear() + w = _weight(n, k, seed, exceptions) + _CODED[key] = (w, encode_chunked(w)) + return _CODED[key] + + +def _args(coded, x: torch.Tensor, n: int, grid: int): + packed, base, exc_off, exc_pos, exc_w = coded + out = torch.full((x.shape[0], n), float("nan"), dtype=torch.bfloat16, device="cuda") + part = torch.zeros(grid * CTA_WARPS * 256, dtype=torch.float32, device="cuda") + counter = torch.zeros((n + 15) // 16, dtype=torch.int32, device="cuda") + return (packed, x, out, part, counter, exc_off, exc_pos, exc_w.view(-1), base, n, grid, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The kernel against its reference and the in-tree BF16 GEMM: the error against fp64 is the in-tree GEMM's + (both are one BF16 rounding of an fp32 sum of the same exact products), the argmax is the in-tree one, nearly + every logit is the in-tree one bit for bit, and a second call gives the same bits.""" + spec = cell.spec + m, n, k = spec["m"], spec["n"], spec["k"] + exceptions = spec.get("exceptions", False) + w, coded = _coded(n, k, 1, exceptions) + if exceptions: + assert coded[3].numel() > 0, "the cell must exercise the exception lists" + g = torch.Generator(device="cuda").manual_seed(_seed(cell.why)) + x = torch.randn(m, k, device="cuda", generator=g).to(torch.bfloat16) + args = _args(coded, x, n, _num_sms()) + certifying(op) + ref = op.reference(*args) + out = op(*args).clone() + again = op(*_args(coded, x, n, _num_sms())).clone() + torch.cuda.synchronize() + assert not bool(args[4].any()), "the split-tile counters are not re-armed" + assert torch.equal(out, again), "the split-tile reduction is not deterministic" + stock = torch.nn.functional.linear(x, w) + exact = x.double() @ w.double().t() + # Each side: one BF16 rounding (<= half an ulp of the logit) of an fp32 sum whose order error is far below it. + err_out = (out.double() - exact).abs().max() + err_stock = (stock.double() - exact).abs().max() + err_ref = (ref.double() - exact).abs().max() + assert err_out <= 1.01 * max(err_stock, err_ref) + 1e-6, (err_out, err_stock, err_ref) + assert (out == stock).float().mean() > 0.99 + assert (out == ref).float().mean() > 0.99 + assert torch.equal(out.float().argmax(-1), stock.float().argmax(-1)) + + +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_integer_data_is_the_bf16_head_bit_for_bit(m) -> None: + """Small-integer weights and activations: every fp32 sum is exact in any order, so the coded head's logits are + the in-tree BF16 GEMM's bit for bit (exceptions included: zeros and integers outside the exponent window).""" + g = torch.Generator(device="cuda").manual_seed(_seed(f"integer {m}")) + n, k = 4096, HIDDEN + w = torch.randint(-3, 4, (n, k), generator=g, device="cuda").to(torch.bfloat16) + w[::7, ::5] = 96.0 # far above the window of the {1, 2, 3} exponents: exceptions + x = torch.randint(-2, 3, (m, k), generator=g, device="cuda").to(torch.bfloat16) + coded = encode_chunked(w) + assert coded[3].numel() > 0 + assert round_trip_ok(*coded, w) + with validating(op): + out = op(*_args(coded, x, n, _num_sms())) + assert torch.equal(out, torch.nn.functional.linear(x, w)) + + +def test_the_head_writes_only_the_shards_columns() -> None: + """Integer data over a 4000-row shard: every logit is the in-tree GEMM's bit for bit, and the columns past n of a + wider output are left untouched.""" + n, k, m = 4000, 1024, 2 + g = torch.Generator(device="cuda").manual_seed(_seed("partial last tile")) + # Small integers: every fp32 sum is exact in any order, so the logits are comparable bit for bit. + w = torch.randint(-3, 4, (n, k), generator=g, device="cuda").to(torch.bfloat16) + w[::7, ::5] = 96.0 # exceptions, including in the partial tile + x = torch.randint(-2, 3, (m, k), generator=g, device="cuda").to(torch.bfloat16) + coded = encode_chunked(w) + assert round_trip_ok(*coded, w) + args = list(_args(coded, x, n, _num_sms())) + args[2] = torch.full((m, n + 16), float("nan"), dtype=torch.bfloat16, device="cuda") + with validating(op): + op(*args) + out = args[2] + assert torch.isnan(out[:, n:]).all(), "the head wrote past the shard's last column" + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out[:, :n], stock) + + +def test_a_shard_of_partial_tiles_keeps_the_bf16_head() -> None: + """The kernel streams whole 16-row tiles: a shard whose rows are not a multiple of 16 is not taken (the fused + verify keeps the in-tree GEMM), and the op refuses it.""" + w = torch.zeros(4004, 1024, dtype=torch.bfloat16, device="cuda") + assert not target_head_kernels.head_supported(SimpleNamespace(weight=w, bias=None)) + assert target_head_kernels.head_supported(SimpleNamespace(weight=w[:4000], bias=None)) + + +@pytest.mark.parametrize("n,k", [(TARGET_SHARD, HIDDEN), (4004, 1024)]) +def test_round_trip_is_bitwise(n, k) -> None: + w, coded = _coded(n, k, 1, True) + assert coded[3].numel() > 0 + assert round_trip_ok(*coded, w) + assert torch.equal(decode(*coded, n, k).view(torch.int16), w.view(torch.int16)) + + +def test_a_corrupted_coding_fails_closed() -> None: + """A flipped sign/mantissa bit, a wrong exception value or a wrong base each fail the bitwise round trip.""" + w = _weight(4004, 1024, 2, exceptions=True) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + bad = packed.clone() + bad[100] ^= 1 + assert not round_trip_ok(bad, base, exc_off, exc_pos, exc_w, w) + bad_w = exc_w.clone() + bad_w[0, 0] += 1e-3 + assert not round_trip_ok(packed, base, exc_off, exc_pos, bad_w, w) + assert not round_trip_ok(packed, base + 1, exc_off, exc_pos, exc_w, w) + + +def test_prepare_head_keeps_only_a_bitwise_coding(monkeypatch) -> None: + """prepare_head keeps the coded shard when it round-trips and runs it as the entry; a coding that does not + round-trip is dropped (the in-tree GEMM stays the path).""" + w = _weight(4000, 1024, 3, exceptions=True) + head = target_head_kernels.prepare_head(SimpleNamespace(weight=w, bias=None)) + assert head.coded and head.weight is w + x = torch.randn(4, 1024, device="cuda").to(torch.bfloat16) + out = target_head_kernels.logits(head, x) + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out.float().argmax(-1), stock.float().argmax(-1)) + assert (out == stock).float().mean() > 0.99 + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[7] ^= 0x40 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(target_head_kernels, "encode_chunked", corrupt) + dropped = target_head_kernels.prepare_head(SimpleNamespace(weight=w, bias=None)) + assert not dropped.coded and dropped.w is None + + +def test_guard_refuses_a_dirty_counter() -> None: + w, coded = _coded(4000, 1024, 1, True) + args = list(_args(coded, torch.zeros(2, 1024, dtype=torch.bfloat16, device="cuda"), 4000, 8)) + args[4][3] = 1 + with pytest.raises(AssertionError, match="counter must be zero"): + with validating(op): + op(*args) + + +def _tiny_exceptions_head(n: int = 4096, k: int = 1024) -> torch.Tensor: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1: exceptions far + below the window whose exact products are the only non-zero logits.""" + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + return w + + +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_tiny_exception_products_are_exact(m) -> None: + """An exception adds exactly x * w: with a one-hot activation on column 0 the logits are 0, 2^-30 and 2^-29 bit + for bit as the BF16 head's, and row 1 wins (a correction against a non-zero placeholder cancels them to 0).""" + w = _tiny_exceptions_head() + coded = encode_chunked(w) + assert coded[1] == 113 and round_trip_ok(*coded, w) + x = torch.zeros((m, w.shape[1]), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + with validating(op): + out = op(*_args(coded, x, w.shape[0], _num_sms())) + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)) + assert out.float().argmax(-1).tolist() == [1] * m + + +def test_a_shard_smaller_than_the_grid_is_stored() -> None: + """A shard with fewer coded units than warps of the requested grid: the launch narrows the grid so every warp owns + a unit, and every logit is stored.""" + w = _tiny_exceptions_head(n=256) + coded = encode_chunked(w) + x = torch.zeros((2, w.shape[1]), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + out = op(*_args(coded, x, w.shape[0], _num_sms())) + assert torch.equal(out.view(torch.int16), torch.nn.functional.linear(x, w).view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..a04f1da1ed81 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_moe_operands.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One tp4 rank's MiniMax-M3 MoE layer operands for the decode-engine entries' tests. + +The routed experts are random checkpoint-format NVFP4 tensors ([w3 (up); w1 (gate)] rows, e4m3 block scales) put +through TRT-LLM's own TRTLLM-Gen NVFP4 loader shuffle (``NVFP4TRTLLMGenFusedMoEMethod._shuffle_all_experts``), which +writes the buffers the engines read in place. The per-expert scalars follow the loader's conventions (act_beta and +act_clamp divided by fc31_alpha, fc31_scale_c = fc2_input_scale x fc31_alpha) and are sized so the FC1 sums land +around the SwiGLU-OAI clamp and the outputs near unit scale. The shared expert's MXFP8 weights come from +mxfp8_quantize with the swizzled scale layout, as the MXFP8 linear method stores them. +""" + +import functools +from types import SimpleNamespace + +import torch + +H, I_TP, NUM_EXPERTS, TOP_K = 6144, 768, 128, 4 +SHARED_I = 768 +SWIGLU_ALPHA, SWIGLU_LIMIT = 1.702, 7.0 +ROUTED_SCALING = 2.0 +MINIMAX2_ROUTING = 5 # RoutingMethodType.MiniMax2 +# NVFP4 global scales: 448 x 6 over the amax each side is sized for. +FC31_INPUT_SCALE = 448.0 * 6.0 / 8.0 +FC2_INPUT_SCALE = 448.0 * 6.0 / 24.0 + + +def _rand_e4m3_scales(shape, gen) -> torch.Tensor: + """E4M3 block scales in [0.5, 2): biased exponents 6 and 7, any mantissa.""" + return torch.randint(0x30, 0x40, shape, dtype=torch.uint8, device="cuda", generator=gen).view( + torch.float8_e4m3fn + ) + + +@functools.lru_cache(maxsize=None) +def experts(seed: int = 20261013): + """The routed experts and the router of one MoE layer, in the layout and conventions of the TRTLLM-Gen loader. + + Returns a SimpleNamespace: ``tensors`` (the ten expert operands in ``M3MoeState.layer``'s order), ``bias``, + ``router_weight``, ``fc31_input_scale`` and ``shared`` (gate_up, its scales, down, its scales). + """ + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + gen = torch.Generator(device="cuda").manual_seed(seed) + e = NUM_EXPERTS + w31 = torch.randint( + 0, 256, (e, 2 * I_TP, H // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + w2 = torch.randint(0, 256, (e, H, I_TP // 2), dtype=torch.uint8, device="cuda", generator=gen) + w31s = _rand_e4m3_scales((e, 2 * I_TP, H // 16), gen) + w2s = _rand_e4m3_scales((e, H, I_TP // 16), gen) + module = SimpleNamespace( + w3_w1_weight=w31, + w2_weight=w2, + w3_w1_weight_scale=w31s, + w2_weight_scale=w2s, + is_gated_activation=True, + bias=False, + ) + NVFP4TRTLLMGenFusedMoEMethod()._shuffle_all_experts(module) + torch.cuda.synchronize() + + # Raw FC1 sums: rms(x) x FC31_INPUT_SCALE x rms(E2M1 value) x rms(block scale) x sqrt(H) ~ 1.0e5 at rms(x) = 1; + # fc31_alpha brings them to ~1.5 in the dequantized domain, the FC2 alpha the outputs to ~1. + base = 1.5 / (FC31_INPUT_SCALE * 2.93 * 1.3 * H**0.5) + fc31_alpha = base * (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) + fc2_alpha = (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) / ( + FC2_INPUT_SCALE * 2.93 * 1.3 * I_TP**0.5 + ) + tensors = ( + w31, + w31s, + w2, + w2s, + fc31_alpha, + FC2_INPUT_SCALE * fc31_alpha, # fc31_scale_c + torch.full((e,), SWIGLU_ALPHA, device="cuda"), # act_alpha + 1.0 / fc31_alpha, # act_beta: the (up + 1) offset over fc31_alpha + SWIGLU_LIMIT / fc31_alpha, # act_clamp over fc31_alpha + fc2_alpha, + ) + bias = (0.02 * torch.randn(e, device="cuda", generator=gen)).float() + router_weight = (torch.randn(e, H, device="cuda", generator=gen) / H**0.5).float().contiguous() + fc31_input_scale = torch.tensor([FC31_INPUT_SCALE], dtype=torch.float32, device="cuda") + + gu = (0.02 * torch.randn(2 * SHARED_I, H, device="cuda", generator=gen)).to(torch.bfloat16) + dn = (0.05 * torch.randn(H, SHARED_I, device="cuda", generator=gen)).to(torch.bfloat16) + gu_q, gu_sf = torch.ops.trtllm.mxfp8_quantize(gu, True) + dn_q, dn_sf = torch.ops.trtllm.mxfp8_quantize(dn, True) + shared = (gu_q.view(2 * SHARED_I, H), gu_sf, dn_q.view(H, SHARED_I), dn_sf) + return SimpleNamespace( + tensors=tensors, + bias=bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared=shared, + ) + + +# Routing cases: random (up to 4 M distinct experts), and every token on the same four experts (the bias outweighs +# any sigmoid difference), which leaves later slot groups empty. +CONCENTRATED = (5, 40, 77, 126) + + +def layer_for(state, case: str, shared: bool = True): + """One layer on ``state`` over :func:`experts` for a routing case. ``shared=False``: the shared expert's weights + are zero, so the engine returns the routed part alone (bf16(0 + bf16(routed))).""" + ops = experts() + bias = ops.bias.clone() + if case == "concentrated": + bias[list(CONCENTRATED)] += 1.0 + shared_weights = ops.shared + if not shared: + gu, gu_sf, dn, dn_sf = ops.shared + shared_weights = (torch.zeros_like(gu), gu_sf, torch.zeros_like(dn), dn_sf) + return state.layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=shared_weights, + ) + + +def stock_routed( + hidden: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """The stock TRTLLM-Gen NVFP4 MoE (fp4_quantize, then fp4_block_scale_moe_runner with SwiGLU-OAI) on the given + routing, over :func:`experts`: what the built-in MiniMaxM3MoE computes for the routed part.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + + ops = experts() + w31, w31s, w2, w2s, fc31_alpha, scale_c, act_alpha, act_beta, act_clamp, fc2_alpha = ops.tensors + xq, xsf = torch.ops.trtllm.fp4_quantize(hidden, ops.fc31_input_scale, 16, False, False) + out = torch.ops.trtllm.fp4_block_scale_moe_runner( + None, None, xq, xsf.view(torch.float8_e4m3fn).flatten(), + w31, w31s, None, act_alpha, act_beta, act_clamp, w2, w2s, None, + scale_c, fc31_alpha, fc2_alpha, + NUM_EXPERTS, TOP_K, None, None, I_TP, 0, NUM_EXPERTS, ROUTED_SCALING, MINIMAX2_ROUTING, True, 0, + topk_weights=topk_weights.to(torch.bfloat16), topk_ids=topk_ids.to(torch.int32), + ) # fmt: skip + return out[0] if isinstance(out, (list, tuple)) else out + + +def hidden_states(m: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(m, H, device="cuda", generator=gen).to(torch.bfloat16) + + +def check_routed_part(op, state, m: int, seed: int) -> None: + """Exact, independent of the reference: the engine's routing (its route_dbg record: ids, then the bf16 weights' + fp32 bits) is the reference routing bit for bit, and its routed part is the stock TRTLLM-Gen MoE's output on that + routing bit for bit -- the claim the engines are built on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=False) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + stock = stock_routed(hidden, ids, weights) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)), ( + "routed part differs from TRTLLM-Gen's" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py new file mode 100644 index 000000000000..724a2979c234 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared operands and gates of the router entry's tests (m3_router_tc).""" + +import zlib + +import torch + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__all__ = ["_build", "_compare_logits", "_seed", "tc"] + + +def _build(m: int, e: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(e, k, generator=g, device="cuda") * 0.02 + return x, w + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +def _compare_logits(out: torch.Tensor, ref: torch.Tensor) -> None: + """fp32 logits that differ from the reference only in summation order. + + Over K = 6144 products a reordered fp32 sum moves by a few ulps of the sum of magnitudes, which near + cancellation is ~sqrt(K) times the row's largest logit's ulp at most. Gated in fp32 ulps of each row's largest + logit: 64 element-wise, 8 in relative RMS. The tensor-core logits against a true fp32 product (the TF32 rounding + the band must see) sit at 2200+ / 1700 ulp. + """ + assert out.dtype == ref.dtype == torch.float32 and out.shape == ref.shape + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=8.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py new file mode 100644 index 000000000000..4ad3e870f33a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe16_engine entry: the 9..13-row MoE engine on its caller-owned state. + +One GB300 rank of the tp4 deployment, the target's build (128 routed CTAs + 24 shared-expert CTAs). Operands and the +band: ``_m3_moe_operands`` and ``test_modeling_v2_m3_moe_engine`` (the same math). +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe16_engine import ( # noqa: E501 + M3Moe16State, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe16_engine import ( # noqa: E501 + m3_moe16_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import H, check_routed_part, hidden_states, layer_for # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_moe16_engine requires a CUDA device" + +# As test_modeling_v2_m3_moe_engine.compare (same arithmetic, same quantization boundaries). +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + +_STATES = {} + + +def _state(m: int) -> M3Moe16State: + if m not in _STATES: + _STATES[m] = M3Moe16State(torch.device("cuda"), num_tokens=m) + return _STATES[m] + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + assert_within_ulp(out, op.reference(hidden, layer), element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +@pytest.mark.parametrize("m", (9, 13)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..c77b345c4bc7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe_engine entry: the 1..8-token MoE engine on its caller-owned state. + +One GB300 rank of the tp4 deployment: 128 experts (all local, intermediate 768 per rank), top 4, hidden 6144, the +shared expert folded in, ``num_ctas`` = SMs - 24 as the target builds it. Operands: ``_m3_moe_operands``. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 # noqa: E501 + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import H, check_routed_part, hidden_states, layer_for # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +_STATES = {} + + +def _state(m: int) -> M3MoeState: + if m not in _STATES: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + _STATES[m] = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + return _STATES[m] + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """The engine against the plain-torch MoE: BAND_ELEMENT bf16 ulp of the row's largest magnitude per element, + BAND_RMS ulp relative RMS. + + Both sides quantize the same operands with the same recipes (the activation's and the intermediate's NVFP4 codes + and block scales, the shared expert's MXFP8); they differ in summation order (tensor-core fp32 against fp64) and in + the PTX approximations (rcp / ex2 / tanh). Such a difference changes the output only where it moves a value across + a rounding boundary: one bf16 ulp of a slot output or of the sum, or, rarely, one intermediate NVFP4 code, which + reaches the output as one weight column's share. Measured over every cell, both routing cases and 3 seeds each + (GB300): at most 1.48 ulp element, 0.18 ulp RMS. The band leaves ~3x on both; a misread recipe (the activation's + 1 / (sf x (1 / sc)) spelling of the reciprocal on ~0.06% of its codes) measured 12 ulp element, 7 ulp RMS. + """ + assert_within_ulp(out, ref, element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + # The workspace is tagged per call and never reset: a repeated call is the same call. + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + compare(out, op.reference(hidden, layer)) + + +@pytest.mark.parametrize("m", (1, 2, 8)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +_GB_STATES = {} + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("m", range(1, 9)) +def test_engine_is_the_gb_tp4_engine_bit_for_bit(m, case) -> None: + """The target's engine (the Rubin m3_moe at the TP4 widths) against the GB TP4 engine it replaces, same operands + and CTA count: the per-expert accumulation order is the same, so the outputs are equal bit for bit.""" + # The shared GB300 engine leaves the tree once no target imports it. + gb_op = pytest.importorskip("tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.moe.op") + if m not in _GB_STATES: + sms = torch.cuda.get_device_properties(0).multi_processor_count + _GB_STATES[m] = gb_op.M3MoeState( + torch.device("cuda"), num_tokens=m, num_ctas=sms - 24, shared=True + ) + hidden = hidden_states(m, seed=300 + m) + out = layer_for(_state(m), case)(hidden) + ref = layer_for(_GB_STATES[m], case)(hidden) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)) + + +def test_layers_on_one_state_do_not_leak_into_each_other() -> None: + """Two layers share one state's workspace, in stream order: A, B, A gives A's output twice, bit for bit, and the + caller's ``out`` receives the same rows as a fresh tensor.""" + state = _state(4) + a, b = layer_for(state, "random"), layer_for(state, "concentrated") + hidden = hidden_states(4, seed=7) + first = op(hidden, a) + op(hidden, b) + out = torch.empty(4, H, dtype=torch.bfloat16, device="cuda") + second = op(hidden, a, out=out) + assert second.data_ptr() == out.data_ptr() + assert torch.equal(first.view(torch.int16), out.view(torch.int16)) + + +@pytest.mark.parametrize( + "bad", + ("misaligned hidden", "fp32 out", "short out", "long out", "int32 arp"), +) +def test_malformed_buffers_are_rejected_by_the_call(bad) -> None: + """Why the entry has no is_valid: the compiled call checks each buffer against its signature and raises.""" + state = _state(2) + layer = layer_for(state, "random") + hidden, kw = hidden_states(2, seed=3), {} + if bad == "misaligned hidden": + hidden = torch.randn(2 * H + 8, device="cuda").to(torch.bfloat16)[2 : 2 + 2 * H].view(2, H) + elif bad == "fp32 out": + kw["out"] = torch.empty(2, H, dtype=torch.float32, device="cuda") + elif bad == "short out": + kw["out"] = torch.empty(1, H, dtype=torch.bfloat16, device="cuda") + elif bad == "long out": + kw["out"] = torch.empty(3, H, dtype=torch.bfloat16, device="cuda") + else: + kw["arp"] = torch.zeros(8, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + op(hidden, layer, **kw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py new file mode 100644 index 000000000000..2a82033abc0b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The TP4 MoE engine's FLAT stack is applied only up to FLAT_MAX_TOKENS (4) tokens per call. + +At 8 tokens FLAT was +2.27 us/call slower than the per-m tables, so above the cap the engine keeps the per-m +configuration (bitwise the same outputs). CPU only: M3MoeState is constructed with its device +queries stubbed and stopped at the kernel build, so the test reads the engine config and the engagement lines +without a GPU. +""" + +import types + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_moe import ( # noqa: E501 + op, +) + + +class _Built(Exception): + """Raised by the stubbed kernel build: the config is final by then.""" + + +def _engine(monkeypatch, num_tokens): + """(config, info lines) of M3MoeState(num_tokens) up to its kernel build.""" + lines = [] + seen = {} + + def build(config, path): + seen["config"] = dict(config) + raise _Built() + + monkeypatch.setattr(op, "is_supported", lambda device=None: (True, "")) + monkeypatch.setattr(op, "_kernel_module", build) + monkeypatch.setattr(op.torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr( + op.torch.cuda, + "get_device_properties", + lambda device: types.SimpleNamespace(multi_processor_count=op.HIDDEN // 32), + ) + from tensorrt_llm.logger import logger + + monkeypatch.setattr(logger, "info", lambda msg, *a, **k: lines.append(str(msg))) + with pytest.raises(_Built): + op.M3MoeState("cuda", num_tokens, num_ctas=op.HIDDEN // 32) + return seen["config"], [line for line in lines if "MiniMax-M3 MoE engine FLAT (TP4)" in line] + + +def _has_flat(config): + return all(config.get(k) == v for k, v in op.FLAT_STACK.items()) + + +@pytest.mark.parametrize("num_tokens", range(1, 9)) +def test_flat_applies_only_up_to_the_cap(monkeypatch, num_tokens): + config, lines = _engine(monkeypatch, num_tokens) + assert op.FLAT_MAX_TOKENS == 4 + if num_tokens <= 4: + assert _has_flat(config), config + assert lines == [ + "MiniMax-M3 MoE engine FLAT (TP4): b2img=0 lg_repl=16 ar_pre=0 late_trigger=0 " + f"tokens={num_tokens}" + ] + else: + # above the cap: the per-m tables + assert config.get("b2img") == op.B2IMG_BY_M[num_tokens - 1] + assert config.get("ar_pre") == op.AR_PRE_BY_M[num_tokens - 1] + assert lines == [ + f"MiniMax-M3 MoE engine FLAT (TP4) off above 4 tokens (per-m tables): tokens={num_tokens}" + ] + + +def test_flat_applies_helper(): + """op.flat_applies: the FLAT stack up to FLAT_MAX_TOKENS tokens, the per-m tables above.""" + assert [op.flat_applies(m) for m in range(1, 9)] == [True] * 4 + [False] * 4 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..e43c57f109a4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_tc entry of the MiniMax-M3 tp4 target: the TF32 tensor-core router of the unfused MoE's +9..16-row pure generation steps (the arithmetic the built-in gate takes at these rows).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import _build, _compare_logits, _seed, tc # noqa: E402 +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + + +@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +def test_tc_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) + with validating(tc): + outs = tc(x, w, None, False) + ref = tc.reference(x, w, None, False) + assert len(outs) == len(ref) == 1 + _compare_logits(outs[0], ref[0]) + + +def test_tc_is_the_built_in_gates_router_call() -> None: + """The target's entry is the op the built-in gate calls at these rows, with the same arguments: the same logits + bit for bit.""" + x, w = _build(14, 128, 6144, 7) + ours = tc(x, w, None, True)[0] + gate = torch.ops.trtllm.minimax_m3_router_tc(x, w, None, True)[0] + assert torch.equal(ours.view(torch.int32), gate.view(torch.int32)) + + +def test_tc_reference_is_tf32_and_the_gate_sees_it() -> None: + """The tensor-core router rounds the weight to TF32: against a true fp32 product its logits sit far outside + the band, so the TF32 rounding in ``reference`` is load-bearing rather than decoration.""" + x, w = _build(16, 128, 6144, 5) + out = tc(x, w, None, False)[0] + with true_fp32_matmul(): + fp32 = x.float() @ w.t() + with pytest.raises(AssertionError): + _compare_logits(out, fp32) + + +def test_guard_refuses_a_misaligned_activation() -> None: + """A row view starting 8 bytes into the buffer keeps a legal row stride but not a legal bulk-copy source.""" + x, w = _build(9, 128, 6144, 3) + flat = torch.empty(9 * 6144 + 4, dtype=torch.bfloat16, device="cuda") + shifted = flat[4:].view(9, 6144) + shifted.copy_(x) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(tc): + tc(shifted, w, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..c207756c1b49 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/m3_embedding_firstnorm entry: the decode embedding gather + the first layer's RMSNorm. + +The table is the model's full bf16 [200064, 6144] embedding shape (random values at the checkpoint's scale), eps the +checkpoint's 1e-6. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + m3_embedding_firstnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" + +H = 6144 +EPS = 1e-6 + + +@functools.lru_cache(maxsize=None) +def _operands(): + g = torch.Generator(device="cuda").manual_seed(20261013) + table = (0.02 * torch.randn(VOCAB, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return table, weight + + +def _buffers(rows): + return (torch.empty(rows, H, dtype=torch.bfloat16, device="cuda"), + torch.empty(rows, H, dtype=torch.bfloat16, device="cuda")) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + rows, id_dtype = cell.spec["rows"], cell.spec["id_dtype"] + table, weight = _operands() + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device="cuda").to(id_dtype) + raw, out = _buffers(rows) + with validating(op): + normed, residual = op(ids, table, weight, EPS, raw, out) + assert normed.data_ptr() == out.data_ptr() and residual.data_ptr() == raw.data_ptr() + ref_normed, ref_raw = op.reference(ids, table, weight, EPS, raw, out) + # The gather is a copy: exact. + assert torch.equal(raw.view(torch.int16), ref_raw.view(torch.int16)) + # The norm: the kernel's fp32 sum of squares (warp butterflies, then across 4 warps) and fast rsqrt against an + # fp64 reference, then x * rstd * w rounded once to bf16. The fp32 side is ~1e-6 relative off, so the two roundings + # disagree only next to a bf16 rounding boundary: one ulp at most per element, a small fraction in aggregate. + # Measured over these cells (3 seeds each, GB300): at most 0.51 ulp element, 0.026 ulp RMS. + assert_within_ulp(out, ref_normed, element_ulp=1.0, rms_ulp=0.1) + + +def test_ids_outside_the_table_give_zero_rows() -> None: + """Documented, not refused: an id outside [0, V) gathers a zero row whose norm is zero.""" + table, weight = _operands() + ids = torch.tensor([-1, 5, VOCAB], dtype=torch.int64, device="cuda") + raw, out = _buffers(3) + op(ids, table, weight, EPS, raw, out) + assert not raw[0].any() and not raw[2].any() and not out[0].any() and not out[2].any() + assert torch.equal(raw[1], table[5]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..27157668c588 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs shared by the tp4 target's spec entry tests: TP4 dimensions, single-device stand-ins for the cross-rank +exchanges, and the comparison bands.""" + +from __future__ import annotations + +import zlib + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) + +#: MiniMax-M3 at TP4: hidden size, vocabulary, the target head's shard and RMSNorm epsilon. +HIDDEN = 6144 +VOCAB = 200064 +WORLD = 4 +TARGET_SHARD = VOCAB // WORLD +EPS = 1e-6 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): + """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" + uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") + keys = torch.zeros(slot_rows, dtype=torch.int64, device="cuda") + counters = torch.tensor([0, call], dtype=torch.int32, device="cuda") + return uc, keys, counters + + +def plant_peers( + uc: torch.Tensor, + counters: torch.Tensor, + own_best: torch.Tensor, + rank: int, + world: int, + shard: int, + g: torch.Generator, + slot_rows: int = SHARD_ROWS, + ties: bool = False, +) -> None: + """Write the other ranks' keys for each row into this rank's slots, as their multicast pushes would. + + Each peer's maximum is the own row maximum shifted by +-0.25 (well clear of the accumulation-order band), so + some rows are won by a peer and some by this rank; with ``ties`` a peer on an even row reports exactly the own + maximum, and the lower global index must win. + """ + if world == 1: + return + m = own_best.shape[0] + half = int(counters[1].item()) & 1 + for r in range(world): + if r == rank: + continue + sign = torch.where(torch.rand(m, generator=g, device="cuda") < 0.5, -1.0, 1.0) + value = own_best.float() + 0.25 * sign + if ties: + value = torch.where(torch.arange(m, device="cuda") % 2 == 0, own_best.float(), value) + index = r * shard + torch.randint(0, shard, (m,), generator=g, device="cuda") + keys = make_keys(value, index) + for t in range(m): + uc[(half * slot_rows + t) * world + r] = keys[t] + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..d31aebfe2610 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_capture_norm entry (Eagle3 capture with fc_norm folded in).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_capture_norm import ( # noqa: E501 + m3_capture_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" + +LAYERS = 3 + + +def _case(spec, seed): + g = gen(seed) + m, slot = spec["m"], spec["slot"] + h = randn((m, HIDDEN), g, 2.0) + r = randn((m, HIDDEN), g, 2.0) if spec["residual"] else None + w = (1.0 + randn((HIDDEN,), g, 0.1).float()).to(torch.bfloat16) + buf = torch.full((m, LAYERS * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + dst = buf[:, slot * HIDDEN : (slot + 1) * HIDDEN] + return (h, r, w, EPS, dst), buf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer + untouched.""" + args, buf = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + slot = cell.spec["slot"] + others = torch.cat([buf[:, : slot * HIDDEN], buf[:, (slot + 1) * HIDDEN :]], dim=1) + assert bool((others == 7.0).all()) + + +def test_guard_refuses_a_strided_weight() -> None: + (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + w2 = torch.stack([w, w], dim=1)[:, 0] + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(h, r, w2, eps, dst) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_normed.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_normed.py new file mode 100644 index 000000000000..e524f094101e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_normed.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU test: a first-draft-inputs reader that skips the drafter's fc_norm (the capture norm is armed) gets the captured +rows unchanged when every captured layer was written with the step's norm, and otherwise a copy with only the raw +layers normalized by their own fc_norm (M3Eagle3Worker._m3_renorm_raw_captures); the per-forward reset clears the +record and the unconsumed fc fork result.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) + +M3Eagle3Worker = modeling.M3Eagle3Worker + +H = 4 +LAYERS = (1, 29, 56) + + +def _worker(normed): + w = object.__new__(M3Eagle3Worker) + w._m3_capture_normed = set(normed) + return w + + +@pytest.fixture(autouse=True) +def _scaling_norm(monkeypatch): + """The drafter's RMSNorm as a stand-in that multiplies by its weight (CPU, no FlashInfer).""" + monkeypatch.setattr(modeling, "flashinfer_rmsnorm", lambda x, w, eps: x * w) + + +def _drafter(): + # fc_norm i scales by (i + 2): enough to tell which blocks were normalized here. + return SimpleNamespace(w={f"fc_norm{i}": i + 2 for i in range(len(LAYERS))}, rms_eps=1e-6) + + +def test_all_captured_layers_normed_returns_the_rows_as_they_lie() -> None: + hs = torch.randn(4, H * len(LAYERS)) + out = _worker(set(LAYERS))._m3_renorm_raw_captures( + hs, SimpleNamespace(layers_to_capture=LAYERS, hidden_size=H), _drafter() + ) + assert out is hs + + +@pytest.mark.parametrize("normed", [set(), {1, 29}, {29, 56}, {1, 56}]) +def test_raw_layers_get_their_own_fc_norm(normed) -> None: + hs = torch.randn(4, H * len(LAYERS)) + ref = hs.clone() + out = _worker(normed)._m3_renorm_raw_captures( + hs, SimpleNamespace(layers_to_capture=LAYERS, hidden_size=H), _drafter() + ) + assert torch.equal(hs, ref) # the capture buffer itself is not modified + for i, layer_id in enumerate(LAYERS): + block, src = out[:, i * H : (i + 1) * H], ref[:, i * H : (i + 1) * H] + assert torch.equal(block, src if layer_id in normed else src * (i + 2)), layer_id + + +def test_arm_resets_the_record_and_the_fork_result() -> None: + w = _worker(set(LAYERS)) + w._m3_fc_hidden = object() + w._m3_fc_early_stream = None # no fork pending + w._m3_join_pending = lambda: None # (tp8) nothing else pending + w.arm_capture_norm( + None, None, None, None + ) # no drafter: resets the per-forward state and returns + assert w._m3_capture_normed == set() + assert w._m3_fc_hidden is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..2caf1da9304d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_chunk_norm entry (the Eagle3 fc_norm over the captured layers).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_chunk_norm import ( # noqa: E501 + m3_chunk_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, c, h = spec["m"], spec["c"], spec["h"] + buf = randn((m, c * h + spec["extra_cols"]), g, 2.0) + src = buf[:, : c * h + spec["extra_cols"]] + w = (1.0 + randn((c, h), g, 0.1).float()).to(torch.bfloat16) + dst = torch.full((m, c * h), 7.0, dtype=torch.bfloat16, device="cuda") + return (src, w, EPS, dst, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + + +def test_guard_refuses_a_strided_weight() -> None: + src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() + wt.copy_(w) + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(src, wt, eps, dst, pdl) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..47d3707f7182 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_input entry (embedding + both RMSNorms + concat in one launch), against its reference and +against the in-tree drafter's modules.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_input import ( # noqa: E501 + m3_draft_input as op, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + VOCAB, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, h = spec["m"], spec["h"] + table = randn((VOCAB, h), g) + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + if spec["invalid_id"]: + ids[0] = -1 + hidden = randn((4 * m if spec["gather"] else m, h), g) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 2 if spec["gather"] else None + w_e = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + w_h = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + out = torch.full((m, 2 * h), 7.0, dtype=torch.bfloat16, device="cuda") + residual = torch.full((m, h), 7.0, dtype=torch.bfloat16, device="cuda") + return (ids, table, hidden, rows, w_e, w_h, EPS, EPS, out, residual, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Normalized halves within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic in a + different summation order, rsqrtf); the gathered residual and the zero embedding of an invalid id exact.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out, residual = op(*args) + ref_out, ref_res = op.reference(*args) + torch.cuda.synchronize() + assert_bf16_close(out, ref_out) + assert torch.equal(residual, ref_res) + if cell.spec["invalid_id"]: + assert not bool(out[0, : cell.spec["h"]].float().any()) + + +def test_guard_refuses_rows_past_hidden() -> None: + gather = next(c for c in op.CELLS if c.spec["gather"]) + args = list(_case(gather.spec, 1)) + args[3] = args[3] + 100 + with pytest.raises(AssertionError, match="rows must index hidden"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("m,gather", [(4, False), (32, False), (8, True)], ids=str) +def test_the_in_tree_draft_input(m, gather) -> None: + """Against the in-tree drafter's own modules: the Embedding lookup, then the layer's input_layernorm and + hidden_norm (in-tree RMSNorm) and the concat of Eagle3DecoderLayer.forward. Within one bf16 ulp on at most 0.1% of + the elements (the RMS statistic summed in another order); the residual (the carried hidden rows) exact.""" + g = gen(seed_of(f"in-tree draft input {m} {gather}")) + h = HIDDEN + embed = Embedding(VOCAB, h, dtype=torch.bfloat16).cuda() + embed.weight.data = randn((VOCAB, h), g) + norms = [] + for _ in range(2): + norm = RMSNorm(hidden_size=h, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + norms.append(norm) + input_layernorm, hidden_norm = norms + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + hidden = randn((4 * m if gather else m, h), g, 2.0) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 3 if gather else None + carried = hidden[rows] if gather else hidden + stock = torch.cat([input_layernorm(embed(ids)), hidden_norm(carried)], dim=-1) + out, residual = m3_spec_tail.draft_input( + ids, embed.weight, hidden, rows, input_layernorm.weight, EPS, hidden_norm.weight, EPS + ) + torch.cuda.synchronize() + assert_bf16_close(out, stock) + assert torch.equal(residual, carried) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py new file mode 100644 index 000000000000..24987c779549 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter boundary push bookkeeping (kernels/m3_draft_lowm arm_push / take_pushed, dispatch.DRAFT_AR_PUSH): the +boundary all-reduce runs prepushed only when the armed call pushed, and only on that call's own output; any other +input raises instead of reducing rows the exchange does not hold. Host state only (no GPU).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_draft_lowm, +) + + +@pytest.fixture(autouse=True) +def _clean(): + m3_draft_lowm._PUSH_ARM, m3_draft_lowm._PUSHED = None, 0 + yield + m3_draft_lowm._PUSH_ARM, m3_draft_lowm._PUSHED = None, 0 + + +def test_not_pushed_is_false_and_disarms() -> None: + m3_draft_lowm._PUSH_ARM = (1, 2, 3, 0, 8, 8) + assert m3_draft_lowm.take_pushed(torch.zeros(4)) is False + assert m3_draft_lowm._PUSH_ARM is None + + +def test_pushed_output_is_true_once() -> None: + out = torch.zeros(4, 8) + m3_draft_lowm._PUSHED = out.data_ptr() + assert m3_draft_lowm.take_pushed(out) is True + assert m3_draft_lowm._PUSHED == 0 + assert m3_draft_lowm.take_pushed(out) is False + + +def test_other_input_raises() -> None: + out = torch.zeros(4, 8) + other = out.clone() + m3_draft_lowm._PUSHED = out.data_ptr() + with pytest.raises(RuntimeError, match="not the boundary all-reduce's input"): + m3_draft_lowm.take_pushed(other) + assert m3_draft_lowm._PUSHED == 0 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py new file mode 100644 index 000000000000..6d1183397871 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_silu entry (the drafter MLP's plain SwiGLU as one C++ launch): bit for bit the Triton +silu_and_mul at its defaults, over every bf16 gate and every bf16 up pattern.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_silu import ( # noqa: E501 + m3_draft_silu as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" + +ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) + + +def _rows(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """[rows, 2d] gate | up from equally long flat gate and up values, d = 4608 (the TP4 drafter shard).""" + d = 4608 + n = gate.numel() + rows = -(-n // d) + pad = rows * d - n + g = torch.cat([gate, gate[:pad]]).view(rows, d) + u = torch.cat([up, up[:pad]]).view(rows, d) + return torch.cat([g, u], dim=1).contiguous() + + +def _assert_bitwise(x: torch.Tensor, what: str) -> None: + got = op(x) + assert got is not None, what + ref = op.reference(x) + torch.cuda.synchronize() + diff = got.view(torch.int16) != ref.view(torch.int16) + assert not bool(diff.any()), f"{what}: {int(diff.sum())} of {diff.numel()} differ" + + +@pytest.mark.parametrize("up", [0.0, 1.0, -1.0, 3.5, -0.0, 1e-3], ids=lambda v: f"up={v}") +def test_every_gate_pattern_is_bitwise_the_triton_kernel(up) -> None: + gate = ALL.cuda() + _assert_bitwise(_rows(gate, torch.full_like(gate, up)), f"gate sweep, up {up}") + + +@pytest.mark.parametrize("gate", [1.0, -1.0, 0.5, -3.0, 20.0, -0.0], ids=lambda v: f"gate={v}") +def test_every_up_pattern_is_bitwise_the_triton_kernel(gate) -> None: + up = ALL.cuda() + _assert_bitwise(_rows(torch.full_like(up, gate), up), f"up sweep, gate {gate}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed) -> None: + g = gen(100 + seed) + n = 1 << 22 + gate = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + up = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + _assert_bitwise(_rows(gate.view(torch.bfloat16), up.view(torch.bfloat16)), f"seed {seed}") + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The drafter's gate_up output range (and past it), in a CUDA graph as the drafter runs it.""" + g = gen(seed_of(cell.why)) + rows, d = cell.spec["rows"], cell.spec["d"] + x = (4.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :6] = torch.tensor([60.0, -60.0, 0.0, -0.0, 88.0, -88.0], dtype=torch.bfloat16) + certifying(op) + ref = op.reference(x) + eager = op(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(x) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(eager.view(torch.int16), ref.view(torch.int16)) + assert torch.equal(captured.view(torch.int16), ref.view(torch.int16)) + + +def test_declines_what_it_does_not_cover() -> None: + x = torch.zeros(2, 2 * 4608, dtype=torch.bfloat16, device="cuda") + assert op(x[:, 8:]) is None # strided rows + assert ( + op(torch.zeros(2, 2 * 12, dtype=torch.bfloat16, device="cuda")) is None + ) # d not a multiple of 8 + assert op(x.float()) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..50c75941815d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,694 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_step entry (one fused greedy draft pass: the BF16 draft head, streamed from its lossless +12-bit coding or as BF16 fragments, the cross-rank argmax and the draft loop's bookkeeping), and its tokens against the +in-tree drafter's head and greedy draft sampling. + +The in-tree path: each rank's LM head shard GEMM (BF16 out), ``LogitsProcessor``'s float cast, then either the argmax +of the gathered logits (``torch.argmax``) or ``greedy_sample_draft_with_tp_gather``'s per-rank ``torch.max`` and the +pick among the ranks' (index, value) pairs. The edge cases use integer data, whose fp32 sums are exact in any order, +so the kernel and the in-tree path see the same BF16 logits and must agree bit for bit; every in-tree comparison runs +with the coded head (the default) and with the BF16 fragments, and on integer data the coded head's winning keys are +the BF16 fragments' bit for bit. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.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_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_LEN, + DRAFT_SHARD, + HIDDEN, + MAX_TOKENS, + bf16_fragments, + from_fragments, + head_weight, + shard_logits, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_step import ( # noqa: E501 + m3_draft_step as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + WORLD, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" + +T = DRAFT_LEN + 1 +_LOW32 = 0xFFFFFFFF + +# ------------------------------------------------------------------ the per-row rule against the in-tree token +# +# The kernel and the in-tree head (a cuBLAS BF16 GEMM) both round an FP32 sum of the same exact BF16 products to BF16, +# in different summation orders. A row's token can differ from the in-tree one only where some candidate's FP32 sum +# can land on either side of a BF16 rounding midpoint. Each row is held to: +# +# kernel token == in-tree token, or the row is provably undecidable: with the fp64 reference logit L_i (exact BF16 +# products, fp64 sum) and the FP32-order bound E_i = gamma_K * sum_k |w_ik * x_k|, the BF16 value of any FP32 sum of +# candidate i lies in [lo_i, hi_i] = [bf16(L_i - E_i), bf16(L_i + E_i)]; the kernel's token must win for some +# choice of values within those intervals (larger value, then the lower global id), and a BF16 rounding midpoint +# must lie within E of the kernel's or the in-tree token's L. +# +# gamma_K = K * 2^-22 is la-xh2's bound (4x the round-to-nearest unit roundoff 2^-24 per addition: any summation +# order, and a truncating tensor-core accumulator); K counts the row's non-zero products (adding an exact zero does +# not round). lo / hi take one more FP32 step outward before the BF16 rounding, which covers rounding L -+ E to FP32. +# Each undecidable row is reported as a warning with its midpoint distances, and each test reports its row count; any +# other mismatch fails. +GAMMA_UNIT = 2.0**-22 + + +def _bf16_outward(v: torch.Tensor, direction: float) -> torch.Tensor: + f = v.float() + f = torch.nextafter(f, torch.full_like(f, direction)) + b = f.to(torch.bfloat16).float() + return torch.where(b == 0, torch.zeros_like(b), b) + + +def bounds(h: torch.Tensor, w: torch.Tensor): + """(fp64 logits L, bound E, lo, hi) [rows, N] of the head input rows h against the head rows w.""" + cols = (h != 0).any(dim=0) # zero columns add exact zeros + hd, wd = h[:, cols].double(), w[:, cols].double() + logits = hd @ wd.t() + bound = (h != 0).sum(dim=1, keepdim=True).double() * GAMMA_UNIT * (hd.abs() @ wd.abs().t()) + del wd + inf = float("inf") + return logits, bound, _bf16_outward(logits - bound, -inf), _bf16_outward(logits + bound, inf) + + +def midpoint_distance(v: float) -> float: + """Distance of v to the nearest BF16 rounding midpoint.""" + b = torch.tensor(v, dtype=torch.float64).float().to(torch.bfloat16).float() + if float(b) == 0.0: + return abs(v) + ulp = 2.0 ** (int(torch.frexp(b).exponent) - 8) + return min(abs(v - (float(b) + ulp / 2)), abs(v - (float(b) - ulp / 2))) + + +def _winnable(tok: int, ids: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Whether candidate ``tok`` wins the row for some values within the candidates' [lo, hi].""" + k = int((ids == tok).nonzero()[0, 0]) + others = ids != tok + beat = (hi[k] > lo[others]) | ((hi[k] == lo[others]) & (tok < ids[others])) + return bool(beat.all()) + + +def check_row(name: str, kernel_tok: int, stock_tok: int, ids, logits, bound, lo, hi) -> bool: + """The per-row rule for one row over the candidates ``ids`` [C] (logits / bound / lo / hi [C]); True when the + row is undecidable (reported as a warning), False when the tokens match; a mismatch outside the bound fails.""" + if kernel_tok == stock_tok: + return False + assert _winnable(kernel_tok, ids, lo, hi), ( + f"{name}: token {kernel_tok} differs from the in-tree {stock_tok} outside the FP32-order bound" + ) + ik = int((ids == kernel_tok).nonzero()[0, 0]) + i_stock = int((ids == stock_tok).nonzero()[0, 0]) + lk, ls, ek, es = ( + float(logits[ik]), + float(logits[i_stock]), + float(bound[ik]), + float(bound[i_stock]), + ) + dk, ds = midpoint_distance(lk), midpoint_distance(ls) + assert dk <= ek or ds <= es, f"{name}: neither logit is within its bound of a BF16 midpoint" + warnings.warn( + f"[draft-step undecidable] {name}: kernel {kernel_tok} (L {lk:.9g}, midpoint distance {dk:.3e}, bound " + f"{ek:.3e}) vs in-tree {stock_tok} (L {ls:.9g}, midpoint distance {ds:.3e}, bound {es:.3e})" + ) + return True + + +def report(name: str, rows: int, undecidable: int) -> None: + warnings.warn(f"[draft-step rows] {name}: {rows} rows, {undecidable} undecidable") + + +def key_value(keys: torch.Tensor) -> torch.Tensor: + """The fp32 value a key orders by (the inverse of the kernel's order-preserving image).""" + u = (keys >> 32) & _LOW32 + bits = torch.where(u >= 0x80000000, u - 0x80000000, _LOW32 - u) + bits = torch.where(bits >= 2**31, bits - 2**32, bits) # as a signed 32-bit pattern + return bits.to(torch.int32).view(torch.float32) + + +def key_index(keys: torch.Tensor) -> torch.Tensor: + return _LOW32 - (keys & _LOW32) + + +def _bookkeeping(nreq: int, step: int, n_x_rows: int, g: torch.Generator): + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if step == 0 + else None + ) + seq_lens = torch.full((nreq,), T, dtype=torch.int32, device="cuda") if step == 0 else None + next_new = torch.full((nreq, T), -5, dtype=torch.int32, device="cuda") + pos_src = torch.randint(100, 3000, (n_x_rows,), generator=g, device="cuda", dtype=torch.int32) + kv_lens = torch.randint(200, 3000, (nreq,), generator=g, device="cuda", dtype=torch.int32) + ids_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + pos_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + return num_acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next + + +def _head(w: torch.Tensor, coded: bool): + """(wf, (exc_off, exc_pos, exc_w, exp_base)) of a BF16 shard: its 12-bit coding (checked to round-trip) or its + BF16 fragments.""" + if not coded: + return bf16_fragments(w), (None, None, None, 0) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + return packed, (exc_off, exc_pos, exc_w.view(-1), base) + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=(None, None, None, 0), +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + coded head's (exc_off, exc_pos, exc_w, exp_base).""" + acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next = _bookkeeping( + nreq, step, x.shape[0], g + ) + if num_acc is None: + num_acc = acc + if uc is None: + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + out_keys = torch.zeros(nreq, dtype=torch.int64, device="cuda") + tiles = wf.shape[0] if coded[0] is None else coded[0].numel() - 1 + return (x, x_rows, nreq, wf, rank * tiles * 16, keys, counters, uc, uc.data_ptr(), rank, world, + next_new[:, 1:], step, ids_next, pos_src, pos_next, kv_lens, num_acc, seq_lens, DRAFT_LEN, num_sms(), + True, out_keys, *coded) # fmt: skip + + +def _case(spec, seed): + g = gen(seed) + nreq, step, rank, world, vs, k = ( + spec["nreq"], + spec["step"], + spec["rank"], + spec["world"], + spec["vs"], + spec["k"], + ) + pass0 = step == 0 + x = randn((nreq * T if pass0 else nreq, k), g) + wf, coded = _head(randn((vs, k), g, 0.02), spec["coded"]) + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if pass0 + else None + ) + x_rows = (torch.arange(nreq, device="cuda") * T + num_acc.long() - 1) if pass0 else None + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + h = x[x_rows] if pass0 else x + own = shard_logits(h, wf, *coded).max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=MAX_TOKENS, ties=spec.get("ties", False) + ) + return _args(x, x_rows, nreq, wf, rank, world, step, g, uc, keys, counters, num_acc, coded) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The bookkeeping is exact on every row (next positions, KV and sequence lengths). Each row's token is the + in-tree one (this rank's BF16 shard logits from the head's GEMM, their first maximum, against the peers' keys) or + the row is provably undecidable (check_row); the winning key's index is the token, and its value lies in the + token's [lo, hi] (or is the peer's pushed key when a peer wins).""" + args = _case(cell.spec, seed_of(cell.why)) + x, x_rows, nreq, wf, offset, _, counters, uc = args[:8] + rank, world = args[9], args[10] + half = int(counters[1].item()) & 1 + pushed = ( + uc[half * MAX_TOKENS * world : (half * MAX_TOKENS + nreq) * world].view(nreq, world).clone() + ) + h = x[x_rows] if x_rows is not None else x[:nreq] + w = head_weight(wf, x.shape[1], *args[23:27]) + best, idx = local_argmax(torch.nn.functional.linear(h, w).float(), offset) + stock = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, MAX_TOKENS) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert args[6].tolist() == [0, 1] and not bool(args[5].any()), "counters / keys not re-armed" + step = args[12] + others = [c for c in range(args[11].shape[1]) if c != step] + assert torch.equal(args[11][:, others], ref["draft_out"][:, others]) + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + assert torch.equal(args[15], ref["pos_next"]) + assert torch.equal(args[16], ref["kv_lens"]) + if args[18] is not None: + assert torch.equal(args[18], ref["seq_lens"]) + # Candidates per row: this rank's head rows (fp64 bounds) and each peer's pushed key (an exact value). + logits, bound, lo, hi = bounds(h, w) + own_ids = offset + torch.arange(w.shape[0], device="cuda") + peers = [p for p in range(world) if p != rank] + undecidable = 0 + for r in range(nreq): + pv = key_value(pushed[r, peers]).double() + ids = torch.cat([own_ids, key_index(pushed[r, peers])]) + undecidable += check_row( + f"{cell.why[:40]} row {r}", + int(args[13][r]), + int(stock[r]), + ids, + torch.cat([logits[r], pv]), + torch.cat([bound[r], torch.zeros_like(pv)]), + torch.cat([lo[r], pv.float()]), + torch.cat([hi[r], pv.float()]), + ) + report(cell.why[:40], nreq, undecidable) + out_keys = args[22] + assert torch.equal(key_index(out_keys), args[13].long()), "the winning key's index is the token" + got = key_value(out_keys) + for r in range(nreq): + c = int(args[13][r]) - offset + if 0 <= c < w.shape[0]: + assert float(lo[r, c]) <= float(got[r]) <= float(hi[r, c]), ( + f"row {r}: winning value outside [lo, hi]" + ) + else: + assert int(out_keys[r]) in pushed[r].tolist(), ( + f"row {r}: a peer won with a key it did not push" + ) + + +def test_fragments_round_trip() -> None: + g = gen(seed_of("fragments")) + w = randn((64, 96), g) + wf = bf16_fragments(w) + assert wf.shape == (4, 6, 256) + assert torch.equal(from_fragments(wf), w) + head = SimpleNamespace(weight=w, bias=None) + # K 96: never coded, the BF16 fragments. + prepared = draft_step_kernels.prepare_head(head, 3, True) + assert not prepared.coded + assert torch.equal(prepared.wf, wf) and prepared.vocab_offset == 3 * 64 and prepared.weight is w + + +# ------------------------------------------------------------------ against the in-tree drafter, bit for bit + + +def _stock_tokens(shards: list, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The in-tree greedy draft tokens of the rows x over the vocab-sharded BF16 head: (argmax of the gathered float + logits, greedy_sample_draft_with_tp_gather's per-rank (index, value) pick).""" + logits = [torch.nn.functional.linear(x, w).float() for w in shards] + gathered = torch.argmax(torch.cat(logits, dim=-1), dim=-1).to(torch.int32) + combined = [] + for rank, shard_logit in enumerate(logits): + worker = SimpleNamespace(mapping=SimpleNamespace(tp_rank=rank)) + combined.append(Eagle3OneModelWorker._get_local_max_and_combined(worker, shard_logit)) + pick = Eagle3OneModelWorker._get_draft_tokens_from_gathered + pick = getattr(pick, "_torchdynamo_orig_callable", pick) + tp_gather = pick(None, torch.cat(combined, dim=-1)).to(torch.int32) + return gathered, tp_gather + + +def _fused_tokens(shards: list, x: torch.Tensor, coded: bool, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's four-rank pass with + its peers' keys in its slots, as their multicast pushes would put them.""" + g = gen(seed_of("fused tokens")) + nreq = x.shape[0] + heads = [_head(w, coded) for w in shards] + own_keys = [] + for rank, (wf, cd) in enumerate(heads): + args = list(_args(x, None, nreq, wf, rank, 1, step, g, coded=cd)) + args[9] = 0 # one-rank run: this rank's slot only, global ids of `rank` + op(*args) + own_keys.append(args[22].clone()) + tokens = [] + for rank, (wf, cd) in enumerate(heads): + uc, keys, counters = head_workspace(WORLD, slot_rows=MAX_TOKENS) + for peer in range(WORLD): + if peer != rank: + for t in range(nreq): + uc[t * WORLD + peer] = own_keys[peer][t] + args = _args(x, None, nreq, wf, rank, WORLD, step, g, uc, keys, counters, coded=cd) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + tokens.append(args[13].clone()) + return tokens + + +def _integer_head(g: torch.Generator, rows: int = WORLD * DRAFT_SHARD) -> torch.Tensor: + """A [rows, 6144] head of small integers over the first 4 columns (zero elsewhere): its logits against integer + rows are exact integers in any summation order, many of which share a BF16 value.""" + w = torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, :4] = torch.randint(-60, 61, (rows, 4), generator=g, device="cuda").to(torch.bfloat16) + return w + + +def _shards(w: torch.Tensor) -> list: + return [w[r * DRAFT_SHARD : (r + 1) * DRAFT_SHARD].contiguous() for r in range(WORLD)] + + +def _assert_stock(w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None) -> None: + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + if expect is not None: + assert gathered.tolist() == expect, (gathered.tolist(), expect) + for rank, tokens in enumerate(_fused_tokens(shards, x, coded)): + assert torch.equal(tokens, gathered), ( + f"rank {rank}: {tokens.tolist()} vs stock {gathered.tolist()}" + ) + + +def _ones_rows(nreq: int) -> torch.Tensor: + x = torch.zeros((nreq, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[:, :4] = 1.0 + return x + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_forced_bf16_ties_take_the_lowest_global_id(coded) -> None: + """A head row copied onto every rank that wins every row of ones (logit 244, every other at most 240): the lowest + global id must win, inside a shard and across ranks. Then integer rows in [-2, 2], whose logits (up to 480 in + magnitude, BF16 spacing 2 above 256) collide often, exactly or after BF16 rounding.""" + g = gen(seed_of("forced ties")) + w = _integer_head(g) + win = torch.tensor([61.0, 61.0, 61.0, 61.0], device="cuda").to(torch.bfloat16) + for gid in (3 * DRAFT_SHARD + 5, DRAFT_SHARD + 77, DRAFT_SHARD + 9000, 2 * DRAFT_SHARD + 1): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), coded, expect=[DRAFT_SHARD + 77] * 5) + # Random integer rows: ties after BF16 rounding happen naturally. + x = torch.randint(-2, 3, (8, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + x[:, 4:] = 0 + _assert_stock(w, x, coded) + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_signed_zero_maxima_take_the_lowest_global_id(coded) -> None: + """Rows whose maximum is zero: an all-zero head input (every logit +-0) and logits <= 0 with zeros by cancellation + on several ranks. -0 and +0 compare equal, so the lowest global id reaching zero wins.""" + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 0] = -1.0 + w[:, 1] = -1.0 + for gid in (2 * DRAFT_SHARD + 5, DRAFT_SHARD + 40000, 3 * DRAFT_SHARD + 7, 49999): + w[gid, 0] = 1.0 + w[gid, 1] = -1.0 + x = torch.zeros((3, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[1, :2] = 1.0 + x[2, :2] = 2.0 + x[2, 2] = -1.0 # w[:, 2] is zero: a -0 product on every row + _assert_stock(w, x, coded, expect=[0, 49999, 49999]) + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + (1000, 1001, "lower"), # 1001 rounds to 1000 in BF16: a tie, the lower id wins + (1001, 1002, "lower"), # 1001 and the tie-to-even midpoint 1002 both round to 1000 + (1000, 1004, "higher"), # adjacent BF16 values: the larger wins + (1003, 1000, "lower"), # 1003 rounds up to 1004 > 1000 + ], + ids=["one_fp32_step", "midpoint", "adjacent_bf16", "rounds_up"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where, coded) -> None: + """Top-2 logits one to four fp32 steps apart: the winner is the one the BF16 logits decide (the in-tree logits' + dtype), not the fp32 sum. The pair sits inside rank 3's shard or on ranks 1 and 2; every other logit is at most + 240 in magnitude.""" + g = gen(seed_of(f"near top2 {lower} {higher} {where}")) + w = _integer_head(g) + if where == "in_shard": + lo_id, hi_id = 3 * DRAFT_SHARD + 100, 3 * DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD + 49000, 2 * DRAFT_SHARD + 3 + for gid, value in ((lo_id, lower), (hi_id, higher)): + q, r = divmod(value, 4) # four integers <= 251, exact in BF16 + w[gid, :4] = torch.tensor([q + (i < r) for i in range(4)], device="cuda").to(torch.bfloat16) + expect = lo_id if winner == "lower" else hi_id + _assert_stock(w, _ones_rows(2), coded, expect=[expect, expect]) + + +# ------------------------------------------------------------------ non-integer rows at a BF16 rounding midpoint + +#: A BF16 rounding midpoint (between 6.0 and 6.03125; a tie there rounds to the even 6.0) and the FP32 spacing there. +MIDPOINT = 6.0 + 2.0**-6 +FP32_ULP = 2.0**-21 + + +def _near_midpoint_head(g: torch.Generator, targets: dict): + """A full-vocabulary head [4 * 50016, 6144] and one head-input row x with 65 non-zero entries: 1, 2^-12, 2^-20 + and 2^-28 on columns 0..3 and 61 random non-integer BF16 values. Every head row carries random non-integer BF16 + values on those 61 columns (logits about 0.2 in magnitude); each row of ``targets`` (global id -> fp64 value) also + gets columns 0..3 set so its fp64 logit is the target to within 2^-30 (a coarse value, then three finer + corrections, each BF16-rounded).""" + x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device="cuda") + scales = (1.0, 2.0**-12, 2.0**-20, 2.0**-28) + for col, scale in enumerate(scales): + x[0, col] = scale + x[0, 4:65] = randn((61,), g, 0.5) + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 4:65] = randn((WORLD * DRAFT_SHARD, 61), g, 0.05) + xd = x[0].double() + for gid, target in targets.items(): + rest = float(target - (w[gid].double() * xd).sum()) + for col, scale in enumerate(scales): + v = torch.tensor(rest / scale, dtype=torch.float64).to(torch.bfloat16) + w[gid, col] = v + rest -= float(v) * scale + assert abs(rest) <= 2.0**-30, (gid, rest) + return w, x + + +def _check_fused_against_stock( + name: str, w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None +) -> int: + """Each rank's fused tokens against the in-tree tokens under the per-row rule (full-vocabulary candidates); + ``expect`` (a token per row) additionally pins rows the bound decides. Returns the undecidable rows.""" + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + fused = _fused_tokens(shards, x, coded) + for tokens in fused[1:]: + assert torch.equal(tokens, fused[0]), "the ranks disagree on the global token" + logits, bound, lo, hi = bounds(x, w) + ids = torch.arange(w.shape[0], device="cuda") + undecidable = 0 + for r in range(x.shape[0]): + undecidable += check_row( + f"{name} row {r}", + int(fused[0][r]), + int(gathered[r]), + ids, + logits[r], + bound[r], + lo[r], + hi[r], + ) + if expect is not None: + assert int(gathered[r]) == expect[r] and int(fused[0][r]) == expect[r], ( + name, + int(gathered[r]), + int(fused[0][r]), + expect[r], + ) + others = ids != expect[r] + assert _winnable(expect[r], ids, lo[r], hi[r]), ( + f"{name}: the bound does not let {expect[r]} win" + ) + rivals = [int(i) for i in ids[others][hi[r][others] >= lo[r][expect[r]]]] + assert not any(_winnable(i, ids, lo[r], hi[r]) for i in rivals), ( + f"{name}: the row is not decided" + ) + report(name, x.shape[0], undecidable) + return undecidable + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + ( + -2, + 2, + None, + ), # straddling the midpoint by 2 FP32 ulps: either side is in-tree up to the order + (2, -2, None), + (0, -2, None), # exactly on the midpoint (ties to the even 6.0) against just below it + ( + -2048, + 2048, + "higher", + ), # ~1e-3 either side, well outside the bound: decided by BF16 rounding + ( + -4096, + -2048, + "lower", + ), # both below the midpoint, ~1e-3 apart: a BF16 tie the lower id wins + ], + ids=["straddle", "straddle_reversed", "on_midpoint", "decided_up", "decided_tie"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where, coded) -> None: + """Non-integer rows whose top-2 fp64 logits sit the given number of FP32 ulps from a BF16 rounding midpoint, inside + rank 3's shard or on ranks 1 and 2, every other logit far below. The per-row rule holds on every row; outside the + bound the token is the one the BF16 rounding decides, bit for bit.""" + g = gen(seed_of(f"midpoint {lower} {higher} {where}")) + if where == "in_shard": + lo_id, hi_id = 3 * DRAFT_SHARD + 100, 3 * DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD + 49000, 2 * DRAFT_SHARD + 3 + w, x = _near_midpoint_head( + g, {lo_id: MIDPOINT + lower * FP32_ULP, hi_id: MIDPOINT + higher * FP32_ULP} + ) + expect = None if winner is None else [lo_id if winner == "lower" else hi_id] + _check_fused_against_stock(f"midpoint {where} {lower:+d}/{higher:+d} ulp", w, x, coded, expect) + + +# ------------------------------------------------------------------ the coded head against the BF16 fragments + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_coded_head_is_the_bf16_head_bit_for_bit(nreq) -> None: + """Integer data (fp32 sums exact in any order) with exceptions on both sides of the exponent window: the coded + head's winning keys (value and global id) and its bookkeeping are the BF16 fragments' bit for bit, on one rank and + across the ranks.""" + g = gen(seed_of(f"coded vs bf16 {nreq}")) + sign = torch.randint(0, 2, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * 2 - 1 + w = (torch.randint(1, 41, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * sign).to( + torch.bfloat16 + ) + # Exceptions on both sides of the window of 1..40: exact zeros, and 2^15 (fp32 sums stay below 2^24). + w[::5, ::512] = 0.0 + w[::13, 7::600] = 32768.0 + x = torch.randint(-2, 3, (nreq, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + wf_c, cd = _head(w, True) + assert cd[1].numel() > 0, "the shard must carry exceptions" + wf_b, cb = _head(w, False) + outs = [] + for wf, c in ((wf_c, cd), (wf_b, cb)): + args = _args(x, None, nreq, wf, 2, 1, 1, gen(seed_of("bookkeeping")), coded=c) + args = list(args) + args[9] = 0 + with validating(op): + op(*args) + torch.cuda.synchronize() + outs.append(args) + coded_args, bf16_args = outs + for i in (11, 13, 15, 16, 22): # draft_out, ids_next, pos_next, kv_lens, out_keys + assert torch.equal(coded_args[i], bf16_args[i]), i + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 2 * DRAFT_SHARD) + assert torch.equal(coded_args[22], make_keys(best, idx)), ( + "the winning keys are not the in-tree ones" + ) + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_tiny_exception_products_are_exact(nreq) -> None: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1, and a one-hot + activation on column 0: the exceptions' exact products are the only non-zero logits, so the coded head must pick + row 1 with the BF16 head's key bit for bit (a correction against a non-zero placeholder cancels them to 0).""" + n, k = 256, 1024 + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + x = torch.zeros((nreq, k), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + wf_c, cd = _head(w, True) + assert cd[3] == 113 + keys = [] + for wf, c in ((wf_c, cd), _head(w, False)): + args = _args(x, None, nreq, wf, 0, 1, 1, gen(seed_of("tiny exceptions")), coded=c) + with validating(op): + # 16 tiles on an all-SM grid: every tile is streamed by all 16 warps, which only the fixed-order sum + # adds deterministically (is_valid rejects the atomics there). + op(*args, fixed_order=True) + torch.cuda.synchronize() + assert args[13].tolist() == [1] * nreq + keys.append(args[22]) + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 0) + assert torch.equal(keys[0], keys[1]) and torch.equal(keys[0], make_keys(best, idx)) + + +def test_prepare_head_codes_the_shard_and_fails_closed(monkeypatch) -> None: + """prepare_head keeps the coded shard when it decodes to the weight bit for bit; a coding that does not is + dropped for the BF16 fragments.""" + g = gen(seed_of("prepare coded")) + w = randn((4096, 1024), g, 0.02) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 2, True) + assert prepared.coded and prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(head_weight(prepared.wf, 1024, prepared.exc_off, prepared.exc_pos, prepared.exc_w, + prepared.exp_base), w) # fmt: skip + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[11] ^= 0x10 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(draft_step_kernels, "encode_chunked", corrupt) + dropped = draft_step_kernels.prepare_head(head, 2, True) + assert not dropped.coded and dropped.exc_off is None + assert torch.equal(dropped.wf, bf16_fragments(w)) + + +def test_prepare_head_keeps_bf16_when_coding_is_off(monkeypatch) -> None: + """``coded`` False (``DRAFT_HEAD_CODED`` off) keeps the BF16 fragments even for a shard whose coding would decode + bit for bit.""" + g = gen(seed_of("prepare bf16")) + w = randn((4096, 1024), g, 0.02) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2, False) + assert not prepared.coded and prepared.exc_off is None + assert torch.equal(prepared.wf, bf16_fragments(w)) + + +def test_guard_refuses_head_rows_past_x() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[1] = args[1] + 1000 + with pytest.raises(AssertionError, match="x_rows must index x"): + with validating(op): + op(*args) + + +def test_guard_refuses_too_few_ctas() -> None: + args = list(_case(op.CELLS[1].spec, 2)) + args[20] = 8 + with pytest.raises(AssertionError, match="tiles per CTA"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py new file mode 100644 index 000000000000..e7c9e5c19e6c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The draft step's tile pieces and their summation order (kernels/m3_draft_step: the fixed warp order). + +A CTA's (tile, chunk) range is cut into 16 equal warp segments; the shared-memory atomics add a tile's pieces in +arrival order, which is deterministic only while no tile takes more than 2 pieces. ``fixed_order`` sums them in warp +order instead. The geometry tests enumerate the real grid (the shard on 212 SMs) so a grid or shard change cannot +silently bring back an order-dependent sum; the GPU tests hold the fixed order to the atomics on exact data and to the +warp-order oracle on a planted tile whose sum depends on the order. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( + m3_draft_step as kds, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_SHARD, + MAX_TOKENS, + bf16_fragments, + max_tile_pieces, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, head_workspace # noqa: E402 + +OP = "minimax_m3_sm107tp4_draft_step" +SMS = 212 # Rubin's SMs: the grid the draft step launches on +K = 6144 +CHUNKS = K // 16 +# This target's shard on Rubin: 14-15 tiles per CTA, at most 2 pieces per tile (the atomics are deterministic and the +# fixed order stays off). +ORDER_DEPENDENT = False + +gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + + +@pytest.mark.parametrize("grid", [SMS, kds.GRID]) +def test_real_grid_pieces(grid) -> None: + """One CTA per SM and the launch the target uses (kernels/m3_draft_step.GRID).""" + pieces = max_tile_pieces(DRAFT_SHARD // 16, grid, CHUNKS) + assert (pieces > 2) == ORDER_DEPENDENT, (grid, pieces) + + +def test_ds_grid_launches_the_configured_grid(monkeypatch) -> None: + monkeypatch.setattr(kds, "_num_sms", lambda: SMS) + assert kds.GRID == 209 + assert kds._grid() == min(kds.GRID, SMS) + + +def test_max_tile_pieces_counts_warp_segments() -> None: + # One tile per CTA: all 16 warps stream it. 7 tiles per CTA: a 168-chunk segment can sit inside a tile (TP8). + assert max_tile_pieces(16, 16, CHUNKS) == 16 + assert max_tile_pieces(7, 1, CHUNKS) == 4 + assert max_tile_pieces(8, 1, CHUNKS) == 2 + assert max_tile_pieces(15, 1, CHUNKS) == 2 + + +def test_the_fixed_order_engages_where_the_geometry_needs_it() -> None: + flag = True + x = torch.empty(1, K) + + def head(n_tiles): + return types.SimpleNamespace(wf=torch.empty(n_tiles, 1), exc_off=None, bnd=None) + + assert kds._fixed_order(1, head(DRAFT_SHARD // 16), x, SMS) == (flag and ORDER_DEPENDENT) + assert kds._fixed_order(1, head(DRAFT_SHARD // 16), x, kds.GRID) == (flag and ORDER_DEPENDENT) + assert kds._fixed_order(1, head(16), x, 16) == flag # 16 pieces per tile + + +def _launcher(wf, x, fixed, nctas): + """(launch, ids, keys): one pass's op launch on buffers made here (outside any capture), world 1.""" + uc, keys, counters = head_workspace(1, slot_rows=MAX_TOKENS) + m = x.shape[0] + i32 = dict(dtype=torch.int32, device="cuda") + draft_out = torch.zeros(m, 4, **i32) + ids = torch.zeros(m, **i32) + pos_src, pos_next, kv_lens = ( + torch.full((m,), 100, **i32), + torch.zeros(m, **i32), + torch.full((m,), 500, **i32), + ) + out = torch.zeros(m, dtype=torch.int64, device="cuda") + + def launch(): + getattr(torch.ops.trtllm, OP)(x, None, m, wf, 0, keys, counters, uc, uc.data_ptr(), 0, 1, draft_out, 1, ids, + pos_src, pos_next, kv_lens, None, None, 3, nctas, False, out, + fixed_order=fixed) # fmt: skip + + return launch, ids, out + + +@gpu +@pytest.mark.parametrize("nctas,vs", [(SMS, DRAFT_SHARD), (16, 256)]) +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_fixed_order_equals_the_atomics_on_exact_data(nctas, vs, m) -> None: + """Integer data: every fp32 sum is exact in any order, so both summations give the same bits.""" + g = gen(17 + m) + w = torch.randint(-1, 2, (vs, K), generator=g, device="cuda").to(torch.bfloat16) + x = torch.randint(-1, 2, (m, K), generator=g, device="cuda").to(torch.bfloat16) + wf = bf16_fragments(w) + out = {} + for fixed in (False, True): + launch, ids, keys = _launcher(wf, x, fixed, nctas) + launch() + torch.cuda.synchronize() + out[fixed] = (keys, ids) + assert torch.equal(out[False][0], out[True][0]) + ref = (x.double() @ w.double().t()).float().bfloat16().float().argmax(dim=1).int() + assert torch.equal(out[True][1], ref) + + +@gpu +@pytest.mark.parametrize("graph", [False, True]) +def test_fixed_order_is_the_warp_order_sum_on_a_planted_tile(graph) -> None: + """One tile per CTA (16 pieces): row 0's pieces are +2^25 (warp 0), +1 (warps 1..14), -2^25 (warp 15). In warp + order the +1s are absorbed by 2^25 and the row sums to 0, so the competitor (row 255, logit 0.5) wins; arrival + order could give up to 14. Every launch must give the warp-order result.""" + w = torch.zeros(256, K, device="cuda") + seg = CHUNKS // 16 * 16 # one warp segment in K + w[0, 0] = 2.0**25 + for s in range(1, 15): + w[0, s * seg] = 1.0 + w[0, 15 * seg] = -(2.0**25) + w[255, 0] = 0.5 + wf = bf16_fragments(w.to(torch.bfloat16)) + x = torch.ones(1, K, device="cuda", dtype=torch.bfloat16) + launch, ids, _ = _launcher(wf, x, True, 16) + run = launch + if graph: + launch() + torch.cuda.synchronize() + gr = torch.cuda.CUDAGraph() + cs = torch.cuda.Stream() + with torch.cuda.stream(cs), torch.cuda.graph(gr, stream=cs): + launch() + run = gr.replay + for _ in range(16): + ids.zero_() + run() + torch.cuda.synchronize() + assert int(ids[0]) == 255 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py new file mode 100644 index 000000000000..b865de3646c1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp4 target's speculative step against the in-tree (stock) Eagle3 numerics, at the real TP4 shapes. + +* Verify: the target logits are the in-tree BF16 LM head's own shard GEMM, and the fused shard argmax + acceptance + gives the in-tree greedy acceptance of the gathered logits, bit for bit, on every rank (peers' keys are the real + per-shard maxima, computed by the same kernel). +* Drafter: the worker runs the in-tree drafter's modules and weights; its own drafter arithmetic is the first pass's + fc_norm (chunk_norm, or capture_norm folded into the capture), held here to the in-tree ``apply_eagle3_fc`` RMSNorm + modules within one bf16 ulp (its effect on the fc output and on the full-vocabulary BF16 head's draft tokens is + measured too), and the fused draft loop's draft input and BF16 head, held to the in-tree modules and greedy draft + sampling by their own tests (test_modeling_v2_m3_draft_input.py, test_modeling_v2_m3_draft_step.py). + +Synthetic weights: hidden 6144, vocabulary 200064 (4 x 50016 shards), 3 captured layers, draft length 3. +""" + +import importlib.util +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as shard_accept_op, +) +from tensorrt_llm._torch.models.modeling_speculative import Eagle3ForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + EPS, + HIDDEN, + TARGET_SHARD, + VOCAB, + WORLD, + assert_bf16_close, + craft_drafts, + gen, + head_workspace, + num_sms, + randn, + seed_of, +) + +assert torch.cuda.is_available(), "the exact-numerics checks require a CUDA device" + +DRAFT_LEN = 3 +CAPTURED_LAYERS = 3 +#: Largest share of fc_norm outputs one bf16 ulp off the in-tree RMSNorm (fp32 statistic summed in another order). +NORM_ULP_FRAC = 1e-3 + + +def _stock_accept(logits: torch.Tensor, draft: torch.Tensor): + """The in-tree greedy strict acceptance (``SpecWorkerBase._sample_and_accept_draft_tokens_base``) of the gathered + float logits of a decode-only batch, no forcing.""" + md = SimpleNamespace(is_all_greedy_sample=True, is_ragged_verify=False, is_cuda_graph=False) + worker = SimpleNamespace(force_num_accepted_tokens=0.0) + worker._sample_tokens_for_batch = lambda *a: Eagle3OneModelWorker._sample_tokens_for_batch( + worker, *a + ) + worker._apply_force_accepted_tokens = lambda *a, **k: ( + Eagle3OneModelWorker._apply_force_accepted_tokens(worker, *a, **k) + ) + nreq = draft.shape[0] + return Eagle3OneModelWorker._sample_and_accept_draft_tokens_base( + worker, logits, draft, 0, nreq, md + ) + + +def _lm_head_shard(weight: torch.Tensor, rank: int) -> LMHead: + """Rank ``rank``'s local part of the target's vocab-sharded BF16 LM head at TP4: an LM head over that rank's + 50016 rows (the TP4 module's local GEMM is this module's; its all-gather is the caller's).""" + head = LMHead(num_embeddings=TARGET_SHARD, embedding_dim=HIDDEN, dtype=torch.bfloat16) + head.weight.data = weight[rank * TARGET_SHARD : (rank + 1) * TARGET_SHARD].contiguous() + return head + + +def _run_shard_accept(logits, rank, world, draft, keys_of_peers=None, call=0): + """The fused verify of one rank: its shard logits, the peers' keys planted in its slots.""" + nreq = draft.shape[0] + rows = logits.shape[0] + uc, keys, counters = head_workspace(world, call=call, slot_rows=SHARD_ROWS) + half = call & 1 + for peer, peer_keys in (keys_of_peers or {}).items(): + for t in range(rows): + uc[(half * SHARD_ROWS + t) * world + peer] = peer_keys[t] + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + shard_accept_op( + logits, rank * TARGET_SHARD if world > 1 else 0, keys, counters, uc, uc.data_ptr(), rank, world, + draft, batch_idx, acc, num, gather, 0, next_new, num_sms(), None, + ) # fmt: skip + torch.cuda.synchronize() + return acc, num, gather, next_new + + +def _local_keys(logits: torch.Tensor, rank: int) -> torch.Tensor: + """What rank ``rank`` pushes: the key of its own shard maximum per row, found by the same kernel at world 1.""" + rows = logits.shape[0] + nreq = rows // (DRAFT_LEN + 1) + dummy = torch.full((nreq, DRAFT_LEN), -1, dtype=torch.int32, device="cuda") + acc, _, _, _ = _run_shard_accept(logits, 0, 1, dummy) + local = acc.view(-1).to(torch.int64) + value = logits.float().gather(1, local[:, None]).squeeze(1) + return make_keys(value, local + rank * TARGET_SHARD) + + +@pytest.mark.parametrize( + "nreq,ties", [(1, False), (2, False), (4, False), (8, True), (12, False)], ids=str +) +def test_fused_verify_is_the_stock_greedy_acceptance(nreq, ties) -> None: + """Target logits (in-tree BF16 head) and verify outputs bit-identical to the in-tree path, on all four ranks.""" + g = gen(seed_of(f"verify {nreq} {ties}")) + rows = nreq * (DRAFT_LEN + 1) + x = randn((rows, HIDDEN), g, 1.0) + weight = randn((VOCAB, HIDDEN), g, 0.02) + if ties: + # Two equal head rows on ranks 1 and 3 that win every row: the lower global id must be the token. + win = (x.float().sum(0) / x.float().sum(0).norm() * 4.0).to(torch.bfloat16) + weight[TARGET_SHARD + 77] = win + weight[3 * TARGET_SHARD + 5] = win + heads = [_lm_head_shard(weight, r) for r in range(WORLD)] + # Stock: each rank's LMHead forward, all-gathered, then LogitsProcessor's float cast. + shard_logits = [head(x) for head in heads] + for head, stock in zip(heads, shard_logits): + assert torch.equal(head.apply_linear(x, None), stock), ( + "the verify's shard GEMM is not the head's" + ) + gathered = torch.cat(shard_logits, dim=1).float() + target_tokens = gathered.argmax(dim=-1).to(torch.int32) + if ties: + assert bool((target_tokens == TARGET_SHARD + 77).all()) + draft = craft_drafts(target_tokens, nreq, DRAFT_LEN) + stock_acc, stock_num = _stock_accept(gathered, draft) + peer_keys = {r: _local_keys(shard_logits[r], r) for r in range(WORLD)} + for rank in range(WORLD): + peers = {r: k for r, k in peer_keys.items() if r != rank} + acc, num, gather, next_new = _run_shard_accept( + shard_logits[rank], rank, WORLD, draft, peers + ) + assert torch.equal(acc, stock_acc.to(torch.int32)), f"rank {rank}: accepted tokens" + assert torch.equal(num, stock_num.to(torch.int32)), f"rank {rank}: accepted counts" + expect_gather = torch.arange(nreq, device="cuda") * (DRAFT_LEN + 1) + num.long() - 1 + assert torch.equal(gather, expect_gather), f"rank {rank}: first-pass gather rows" + assert torch.equal(next_new[:, 0], target_tokens[expect_gather]), f"rank {rank}: next_new" + + +class _Drafter: + """The in-tree Eagle3 first-pass pieces at the target's shapes: per-layer fc_norm RMSNorms and the fc.""" + + def __init__(self, g: torch.Generator): + self.norms = [] + for _ in range(CAPTURED_LAYERS): + norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ) + self.norms.append(norm) + self.fc = Linear(HIDDEN * CAPTURED_LAYERS, HIDDEN, bias=False, dtype=torch.bfloat16).cuda() + self.fc.weight.data = randn((HIDDEN, HIDDEN * CAPTURED_LAYERS), g, 0.02) + self.model = SimpleNamespace( + dtype=torch.bfloat16, + hidden_size=HIDDEN, + fc_norm=torch.nn.ModuleList(self.norms), + _norm_before_fc=False, + fc=self.fc, + ) + self.stacked = torch.stack([n.weight.detach() for n in self.norms]).contiguous() + + def stock_normed(self, hs: torch.Tensor) -> torch.Tensor: + """``apply_eagle3_fc``'s fc input: each captured layer through its own in-tree RMSNorm.""" + chunks = hs.chunk(CAPTURED_LAYERS, dim=-1) + return torch.cat([norm(c) for norm, c in zip(self.norms, chunks)], dim=-1) + + def stock_fc(self, hs: torch.Tensor) -> torch.Tensor: + return Eagle3ForCausalLM.apply_eagle3_fc(self, hs) + + +def _captures(rows: int, g: torch.Generator): + """Per captured layer (hidden, residual) and the in-tree capture buffer of bf16(hidden + residual).""" + layers = [(randn((rows, HIDDEN), g, 1.5), randn((rows, HIDDEN), g, 3.0)) for _ in range(3)] + buf = torch.cat([h + r for h, r in layers], dim=-1) + return layers, buf + + +def _row_identical(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a == b).all(dim=-1) + + +@pytest.mark.parametrize("rows", (4, 8, 32, 48), ids=lambda r: f"rows{r}") +def test_first_drafter_inputs_are_the_stock_fc_norm(rows) -> None: + """chunk_norm and capture_norm against ``apply_eagle3_fc``'s in-tree RMSNorms: within one bf16 ulp on at most + NORM_ULP_FRAC of the elements; the fc (the in-tree module in both paths) then differs only on rows whose input + did.""" + g = gen(seed_of(f"fc_norm {rows}")) + drafter = _Drafter(g) + layers, buf = _captures(rows, g) + stock_in = drafter.stock_normed(buf) + stock_out = drafter.stock_fc(buf) + assert torch.equal(drafter.fc(stock_in), stock_out) + + chunk_in = m3_spec_tail.chunk_norm(buf, drafter.stacked, EPS) + capture_in = torch.empty_like(buf) + for i, (h, r) in enumerate(layers): + m3_spec_tail.capture_norm( + h, r, drafter.stacked[i], EPS, capture_in[:, i * HIDDEN : (i + 1) * HIDDEN] + ) + torch.cuda.synchronize() + for name, ours in (("chunk_norm", chunk_in), ("capture_norm", capture_in)): + assert_bf16_close(ours, stock_in, max_frac=NORM_ULP_FRAC) + out = drafter.fc(ours) + same_rows = _row_identical(ours, stock_in) + assert torch.equal(out[same_rows], stock_out[same_rows]), ( + f"{name}: fc moved an unchanged row" + ) + moved = (ours != stock_in).float().mean().item() + print( + f"[exact-numerics] rows={rows} {name}: {moved:.2e} of fc inputs one ulp off, " + f"{int((~same_rows).sum())}/{rows} rows differ, fc max |diff| " + f"{(out.float() - stock_out.float()).abs().max().item():.3e}" + ) + assert torch.equal(chunk_in, capture_in), "capture_norm and chunk_norm disagree" + + +@pytest.mark.parametrize("rows", (4, 32), ids=lambda r: f"rows{r}") +def test_draft_tokens_through_the_full_vocab_head(rows) -> None: + """The fc_norm's ulp-level differences carried through the full-vocabulary BF16 head (a proxy of the draft head + input): the greedy draft token equals the stock one on every row whose stock top-2 margin exceeds the logit + perturbation.""" + g = gen(seed_of(f"draft head {rows}")) + drafter = _Drafter(g) + _, buf = _captures(rows, g) + weight = randn((VOCAB, HIDDEN), g, 0.02) + head = LMHead(num_embeddings=VOCAB, embedding_dim=HIDDEN, dtype=torch.bfloat16) + head.weight.data = weight + stock_logits = head(drafter.stock_fc(buf)).float() + ours_logits = head(drafter.fc(m3_spec_tail.chunk_norm(buf, drafter.stacked, EPS))).float() + perturb = (ours_logits - stock_logits).abs().max(dim=-1).values + top2 = stock_logits.topk(2, dim=-1).values + decidable = (top2[:, 0] - top2[:, 1]) > 2 * perturb + stock_tok, ours_tok = stock_logits.argmax(-1), ours_logits.argmax(-1) + assert torch.equal(ours_tok[decidable], stock_tok[decidable]) + print( + f"[exact-numerics] rows={rows} draft head: {int((ours_tok == stock_tok).sum())}/{rows} tokens equal, " + f"{int(decidable.sum())} decidable, max logit |diff| {perturb.max().item():.3e}" + ) + + +def test_the_drafter_is_the_in_tree_drafter() -> None: + """Outside the fused draft loop's steps the worker drafts with the in-tree loop, forward and sampling, and nothing + re-quantizes the drafter's linears or the LM heads: no hook of a quantized / pruned head or drafter projection is + left in the target, and the fused head streams the BF16 head weight itself.""" + assert ( + modeling.M3Eagle3Worker._forward_draft_loop is not Eagle3OneModelWorker._forward_draft_loop + ) + for name in ( + "_forward_linear_draft_loop", + "_run_draft_forward", + "sample_draft_tokens", + "sample_and_accept_draft_tokens", + ): + assert getattr(modeling.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), name + for module, name in ( + (modeling, "M3DraftGemvLinearMethod"), + (modeling, "adopt_draft_linears"), + (modeling, "greedy_head_outputs"), + (modeling, "M3FP8LMHeadMethod"), + (modeling, "adopt_fp8_lm_heads"), + ): + assert not hasattr(module, name), f"{module.__name__}.{name}" + pkg = modeling.__name__.rsplit(".", 1)[0] + for gone in ("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_107__tp4/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py new file mode 100644 index 000000000000..cf2b42bc86d9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the forward-tail spec fold (kernels/m3_fold_tail and M3SpecStepFold.in_forward). + +The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers +(step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing +when the step was not armed. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as spec_fold, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( + m3_fold_tail, + m3_spec_tail, +) + +assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" + +VOCAB = 200064 +NUM_SLOTS = 9 +MAX_ROWS = 8 +STALE = -4 + + +def _store(widths, g): + wn, wx, wd = widths + i32 = dict(dtype=torch.int32, device="cuda") + return types.SimpleNamespace( + new_tokens=torch.randint( + 0, VOCAB, (wn, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_new_tokens=torch.randint( + 0, VOCAB, (wx, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_draft_tokens=torch.randint( + 0, VOCAB, (NUM_SLOTS, wd), generator=g, device="cuda", dtype=torch.int32 + ), + new_tokens_lens=torch.randint(1, 5, (NUM_SLOTS,), generator=g, **i32), + ) + + +def _clone(ns): + return types.SimpleNamespace(**{k: v.clone() for k, v in vars(ns).items()}) + + +def _outputs(rows, t, g): + d = t - 1 + return { + "new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_draft_tokens": torch.randint( + 0, VOCAB, (rows, d), generator=g, device="cuda", dtype=torch.int32 + ), + "new_tokens_lens": torch.randint( + 1, t + 1, (rows,), generator=g, device="cuda", dtype=torch.int32 + ), + } + + +def _engine(t): + i32 = dict(dtype=torch.int32, device="cuda") + return ( + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS * (t - 1),), STALE, **i32), + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS,), STALE, **i32), + ) + + +def _reference(store, engine, slots, out, t): + """slot_scatter into the store, then step_gather of the next step's inputs from it.""" + assert m3_spec_tail.slot_scatter( + slots, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["new_tokens_lens"], + store.new_tokens_lens, + out["next_draft_tokens"], + store.next_draft_tokens, + ) + p = slots.numel() + ids, draft, pos_off, kv_off = engine + slots32 = slots.to(torch.int32) + assert m3_spec_tail.step_gather( + store.next_new_tokens, + store.next_draft_tokens, + store.new_tokens_lens, + slots32, + slots32.repeat_interleave(t), + p, + t, + t - 1, + ids, + 0, + draft, + 0, + pos_off, + kv_off, + 0, + ) + + +def _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=False): + ids, draft, pos_off, kv_off = engine + assert m3_fold_tail.fold_tail( + ctrl, + slots_buf, + rows, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["next_draft_tokens"], + store.next_draft_tokens, + out["new_tokens_lens"], + store.new_tokens_lens, + ids, + draft, + pos_off, + kv_off, + t, + cpp=cpp, + ) + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("widths", [(4, 4, 3), (8, 8, 7)], ids=["store=step", "store wider"]) +def test_tail_matches_scatter_then_gather(rows, widths, cpp) -> None: + """Store bit for bit; engine rows [0, P*T) / [0, P) bit for bit (the forward reads only those rows).""" + t = 4 + g = torch.Generator(device="cuda").manual_seed(1000 * rows + widths[0]) + store = _store(widths, g) + out = _outputs(rows, t, g) + slots = torch.randperm(NUM_SLOTS, generator=g, device="cuda")[:rows] + ref_store, tail_store = _clone(store), _clone(store) + ref_engine, tail_engine = _engine(t), _engine(t) + _reference(ref_store, ref_engine, slots, out, t) + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + slots_buf = torch.zeros(MAX_ROWS, dtype=torch.long, device="cuda") + slots_buf[:rows] = slots + _tail(tail_store, tail_engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(ref_store, name), getattr(tail_store, name)), name + spans = (rows * t, rows * (t - 1), rows * t, rows) + for k, (ref, got, n) in enumerate(zip(ref_engine, tail_engine, spans)): + assert torch.equal(ref[:n], got[:n]), k + assert torch.all(got[n:] == STALE), k + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +def test_unarmed_tail_writes_nothing(cpp) -> None: + """ctrl[0] = 0 (the step was not armed): the store and the engine buffers are untouched.""" + t, rows = 4, 2 + g = torch.Generator(device="cuda").manual_seed(7) + store = _store((4, 4, 3), g) + out = _outputs(rows, t, g) + before, engine = _clone(store), _engine(t) + ctrl = torch.zeros(1, dtype=torch.int32, device="cuda") + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(before, name), getattr(store, name)), name + for buf in engine: + assert torch.all(buf == STALE) + + +def _fold(): + fold = spec_fold.M3SpecStepFold(in_forward=True) + t = 4 + engine = _engine(t) + fold.bind_engine_buffers(*engine, max_batch_size=MAX_ROWS) + g = torch.Generator(device="cuda").manual_seed(11) + fold.bind_sampler_store(_store((4, 4, 3), g)) + return fold, g, t + + +def test_fold_handshake_stages_only_confirmed_rows() -> None: + """begin_step arms; the sampler confirms only the armed rows from the tail's outputs; the next step is staged only + for exactly those rows in that order.""" + fold, g, t = _fold() + assert fold.begin_step([3, 5], 2) is False # nothing staged before the first forward + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [3, 5]) is True + assert fold.begin_step([3, 5], 2) is True + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert ( + fold.sampler_rows_stored(out, 0, [5, 3]) is False + ) # other order: the sampler stores itself + assert fold.begin_step([5, 3], 2) is False + + +def test_fold_handshake_needs_the_tail_outputs() -> None: + """Outputs the tail did not run on (another dict, a mixed step, an unarmed step) are never reported stored.""" + fold, g, t = _fold() + fold.begin_step([1], 1) + fold.store_in_forward(_outputs(1, t, g)) + assert fold.sampler_rows_stored(_outputs(1, t, g), 0, [1]) is False + fold.begin_step(None, 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [1]) is False + fold.begin_step([1], 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 1, [1]) is False # rows were skipped (chunked context) + + +def test_a_padded_step_after_an_armed_one_writes_nothing() -> None: + """K3 #150 check #5: CUDA-graph padding rows carry no seq slot, so the engine passes rows=None for a padded step. + Even right after an armed step, that step uploads ctrl 0 and its forward tail leaves the store and the engine + buffers untouched (a dummy row never writes a live slot).""" + fold, g, t = _fold() + fold.begin_step([2, 4], 2) + fold.store_in_forward(_outputs(2, t, g)) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 2 + + fold.begin_step(None, 4) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 0 + store_before = _clone(fold.store) + engine_before = [buf.clone() for buf in fold.step_inputs] + fold.store_in_forward(_outputs(4, t, g)) + torch.cuda.synchronize() + for name in vars(store_before): + assert torch.equal(getattr(store_before, name), getattr(fold.store, name)), name + for before, buf in zip(engine_before, fold.step_inputs): + assert torch.equal(before, buf) + + +def test_fold_off_by_default() -> None: + fold = spec_fold.M3SpecStepFold() + assert not fold.in_forward + assert fold.begin_step([0], 1) is False + + +@pytest.mark.parametrize( + "in_forward, store_copy_wait, expected", + [(True, True, True), (True, False, False), (False, True, False)], +) +def test_the_store_copy_event_exists_only_with_the_tail_and_the_wait( + in_forward, store_copy_wait, expected +): + """STEP_SPEC_STORE_COPY_SIDE: an external event the sampler records after its side-stream store copies and the + forward's tail waits on; only a fold that stores in the forward needs it.""" + fold = spec_fold.M3SpecStepFold(in_forward=in_forward, store_copy_wait=store_copy_wait) + device = torch.device("cuda") + fold.bind_engine_buffers( + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + 4, + ) + assert (fold.store_copies_done is not None) == expected + if expected: + fold.store_copies_done.synchronize() + + +def test_the_store_wait_precedes_only_a_plain_launch_under_capture(): + """The store wait is an external-event wait node under graph capture, which drops the next kernel's PDL edge: + it may precede the fold tail (a plain launch); any other launch fails the capture (eager steps are not + checked).""" + fold = spec_fold.M3SpecStepFold(in_forward=True, store_copy_wait=True) + device = torch.device("cuda") + fold.bind_engine_buffers( + *(torch.zeros(8, dtype=torch.int32, device=device) for _ in range(4)), 4 + ) + assert fold.STORE_WAIT_BEFORE == frozenset({"fold_tail"}) + fold._wait_store_copies(before="draft_step") # eager: allowed + x = torch.zeros(4, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + fold._wait_store_copies(before="fold_tail") + x.add_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(x, torch.ones(4, device=device)) + bad = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="PDL edge"): + with torch.cuda.graph(bad, stream=stream): + fold._wait_store_copies(before="draft_step") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..717deed4c1aa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec entry of the MiniMax-M3 tp4 target: the fractional forced-acceptance draw.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(5) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + counter = torch.tensor([41], dtype=torch.int64, device="cuda") + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + counter, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert int(counter.item()) == 42, "the draw bumps the step counter once" + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py new file mode 100644 index 000000000000..8a4d03c2fd03 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec_req entry of the MiniMax-M3 tp4 target: the per-request fractional forced-acceptance +draw (FORCE_AL_PER_REQUEST: one draw per request at its own key, its KV length at the verify, no shared counter), so +a request's forced count does not depend on the call history, on its slot in the batch, or on the other requests.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_force_vec_req import ( # noqa: E501 + m3_force_vec_req as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SEED, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(7) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + # Each request's key is its KV length at the verify: long, distinct, arbitrary. + keys = torch.randint( + 1, 1 << 20, (spec["batch"],), generator=g, device="cuda", dtype=torch.int64 + ).to(torch.int32) + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + keys, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert torch.equal(out, ref) + + +def _pool(): + g = torch.Generator(device="cpu") + g.manual_seed(_FORCE_ACCEPT_RNG_SEED) + return torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, dtype=torch.float32, generator=g).cuda() + + +def _draw(pool, keys, base_total=2, max_total=4, frac=0.6): + out = torch.empty(keys.numel(), dtype=torch.int32, device="cuda") + return op( + pool, + keys, + out, + keys.numel(), + base_total, + max_total, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (keys.numel() - 1).bit_length()), + ) + + +def test_a_request_draws_the_same_whatever_ran_before_or_beside_it() -> None: + pool = _pool() + keys = torch.tensor([165_521, 40_003, 323_007, 7], dtype=torch.int32, device="cuda") + first = _draw(pool, keys) + for _ in range(5): # no shared counter: repeated calls draw the same + assert torch.equal(_draw(pool, keys), first) + perm = torch.tensor([2, 0, 3, 1], device="cuda") + assert torch.equal(_draw(pool, keys[perm]), first[perm]) # slot-independent + alone = torch.cat([_draw(pool, keys[i : i + 1]) for i in range(4)]) + assert torch.equal(alone, first) # independent of the other requests + + +def test_the_counts_follow_the_forced_fraction() -> None: + pool = _pool() + keys = torch.arange(100_000, 100_000 + 64 * 400, dtype=torch.int32, device="cuda") + counts = torch.cat([_draw(pool, keys[i : i + 64]) for i in range(0, keys.numel(), 64)]) + assert set(counts.tolist()) <= {2, 3} + extra = (counts == 3).float().mean().item() + assert abs(extra - 0.6) < 0.02, extra diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py new file mode 100644 index 000000000000..9c7930e2a800 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (kernels/m3_lowm_splitk, +kernels/m3_draft_lowm): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's +linears, eager and in a CUDA graph, with and without PDL.""" + +import pytest +import torch +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + tactic_for, +) + +assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" + +# Every certified (N, K, rows) of the TP4 drafter table that a split-K tactic serves: (N, K, rows, tactic). +CELLS = [ + (nk[0], nk[1], m, tactic_for(nk[0], nk[1], m)[1]) + for nk, rows in sorted(ROWS.items()) + for m in rows + if tactic_for(nk[0], nk[1], m)[0] == "s" +] + + +def test_the_table_has_the_certified_split_k_cells() -> None: + assert sorted((n, k, m) for n, k, m, _ in CELLS) == [ + (2304, 12288, 4), + (6144, 2048, 4), + (6144, 4608, 1), + (6144, 4608, 4), + (9216, 6144, 4), + ] + + +@pytest.mark.parametrize("use_pdl", [False, True], ids=["plain", "pdl"]) +@pytest.mark.parametrize("cell", CELLS, ids=[f"{n}x{k}@{m}" for n, k, m, _ in CELLS]) +def test_bitwise_flashinfer(cell, use_pdl) -> None: + n, k, m, tactic = cell + g = torch.Generator(device="cuda").manual_seed(n + 7 * k + m) + for trial in range(3): + x = (torch.randn(m, k, generator=g, device="cuda") * (1.0 + trial)).to(torch.bfloat16) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + ref = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + got = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + fi_run(x, w.t(), None, ref, use_pdl, FiTactic(*tactic)) + run_splitk_dense(x, w.t(), None, got, use_pdl, SplitKTactic(*tactic)) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), f"trial {trial}" + # In a CUDA graph, as the drafter runs it. + cap = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_splitk_dense(x, w.t(), None, cap, use_pdl, SplitKTactic(*tactic)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(cap.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..51fee0cff1c5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sampled_accept entry (fused temperature + top-p verify) and its workspace size.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_sampled_accept import ( # noqa: E501 + m3_sampled_accept as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_sampled_accept_ws_words import ( # noqa: E501 + m3_sampled_accept_ws_words as ws_words, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" + +DRAFT_LEN = 3 +#: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). +HEAD = 12 +#: Draws per distribution cell, and the total-variation band against the reference distribution. With 12 +#: categories and >= 4000 draws the expected TV of an exact sampler is ~0.02 (sd ~0.006); 0.06 is ~6 sd above it +#: while a wrong nucleus cut (one category too many or too few, >= 1 / 12 of the mass moved) lands at >= 0.08. +DRAWS = 4000 +TV_BAND = 0.06 + + +def _logits(spec, g, rows): + vs = spec["vs"] + if spec["check"] == "flat": + return (torch.randn((rows, vs), generator=g, device="cuda") * 0.3).to(torch.bfloat16) + if spec["check"] == "argmax": + return (torch.randn((rows, vs), generator=g, device="cuda") * 3).to(torch.bfloat16) + row = torch.full((vs,), -20.0, device="cuda") + hot = torch.randperm(vs, generator=g, device="cuda")[:HEAD] + row[hot] = torch.linspace(2.0, 0.2, HEAD, device="cuda") + return row.to(torch.bfloat16).expand(rows, vs).contiguous() + + +def _case(spec, seed): + g = gen(seed) + nreq = spec["nreq"] + rows = nreq * (DRAFT_LEN + 1) + logits = _logits(spec, g, rows) + words = ws_words(1) + uc = torch.zeros(words, dtype=torch.int64, device="cuda") + counters = torch.zeros(2, dtype=torch.int32, device="cuda") + tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + extra = torch.zeros(1, dtype=torch.int32, device="cuda") + temperatures = torch.full((rows,), spec["temperature"], dtype=torch.float32, device="cuda") + top_ps = torch.full((rows,), spec["top_p"], dtype=torch.float32, device="cuda") + seeds = torch.randint(0, 2**62, (rows,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.zeros(rows, dtype=torch.int64, device="cuda") + # Drafts drawn from the likeliest tokens, so some prefixes match. + top = logits[0].float().topk(3).indices.to(torch.int32) + spec["vocab_offset"] + draft = top[torch.randint(0, 3, (nreq, DRAFT_LEN), generator=g, device="cuda")].contiguous() + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + force = spec.get("force", "none") + force_total = 3 if force == "total" else 0 + force_vec = (torch.randint(1, DRAFT_LEN + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if force == "vec" else None) # fmt: skip + return [logits, spec["vocab_offset"], temperatures, top_ps, seeds, offsets, uc, uc.data_ptr(), counters, tokens, + 0, 1, draft, batch_idx, acc, num, gather, force_total, next_new, force_vec, extra, True] # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Statistical for the draw, exact for everything derived from it. + + Every drawn token has nonzero reference probability (lies in the nucleus); the acceptance outputs equal the + reference's strict acceptance of the drawn tokens, every call; the call counter advances once per call; and, + for the distribution cells, the draws over many calls (fresh Philox offsets) are within TV_BAND of the + reference distribution (see DRAWS / TV_BAND). + """ + spec = cell.spec + args = _case(spec, seed_of(cell.why)) + logits, offset, rows = args[0], spec["vocab_offset"], args[0].shape[0] + certifying(op) + ref = op.reference(*args) + probs = ref["probs"] + calls = 1 if spec["check"] == "argmax" else -(-DRAWS // rows) + counts = torch.zeros(logits.shape[1], dtype=torch.float64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + drawn = args[9][:rows].clone() + local = drawn.long() - offset + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + assert bool((probs.gather(1, local[:, None]) > 0).all()), ( + "a drawn token outside the nucleus" + ) + acc = ref["accept"](drawn) + assert torch.equal(args[14], acc["accepted"]) + assert torch.equal(args[15], acc["num_accepted"]) + assert torch.equal(args[16], acc["gather_ids"]) + assert torch.equal(args[18], acc["next_new"]) + assert args[8].tolist() == [0, c + 1] + counts.index_add_(0, local, torch.ones(rows, dtype=torch.float64, device="cuda")) + if spec["check"] == "distribution": + emp = counts / counts.sum() + tv = 0.5 * (emp - probs[0].double()).abs().sum().item() + assert tv <= TV_BAND, f"total variation {tv:.3f} against the reference nucleus distribution" + + +@pytest.mark.parametrize("top_p", [0.001, 0.0002]) +def test_long_rejection_chains_stay_in_the_nucleus(top_p) -> None: + """A tiny nucleus over a flat tail: rows that need more exchange rounds than the workspace has slots. + + Each round's candidate is drawn over the tokens heavier than the last one, so a row takes ~ln(1 / top_p) + 1 + rounds on average, with a tail past the 16 exchange slots (about 1% of the rows at top_p 0.001). Rounds reuse the + slots cyclically, told apart by a lap-parity bit in the wire tags, and no round accepts unconditionally: every + draw, however many rounds it took, lies in the reference nucleus. World 1 (one rank's 8 CTAs per row exchange + through the same slot ring); the cross-rank path needs the multi-device job. + """ + spec = dict(nreq=1, vs=50016, vocab_offset=50016, check="flat", top_p=top_p, temperature=1.0) + args = _case(spec, seed_of(f"flat tail top_p {top_p}")) + logits, rows = args[0], args[0].shape[0] + probs = op.reference(*args)["probs"] + calls = -(-DRAWS // rows) + outside = torch.zeros((), dtype=torch.int64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + local = args[9][:rows].long() - spec["vocab_offset"] + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + outside += (probs.gather(1, local[:, None]) == 0).sum() + assert args[8].tolist() == [0, c + 1] + draws = calls * rows + rounds = 2.0 + args[20].item() / draws + kept = int((probs > 0).sum(1).min().item()) + assert outside.item() == 0, ( + f"{outside.item()} of {draws} draws outside the nucleus (kept >= {kept} of {logits.shape[1]}, " + f"{rounds:.2f} rounds per row on average)" + ) + assert rounds > 4.0, ( + f"only {rounds:.2f} rounds per row: the rows no longer stress the slot ring" + ) + + +@pytest.mark.parametrize("cell", ws_words.CELLS, ids=[c.why[:44] for c in ws_words.CELLS]) +def test_workspace_words(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(ws_words) + world = cell.spec["world"] + assert ws_words(world) == ws_words.reference(world) + + +def test_guard_refuses_top_k_rows() -> None: + args = _case(op.CELLS[0].spec, 1) + top_ks = torch.full((args[0].shape[0],), 20, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="top-k rows"): + with validating(op): + op(*args, top_ks=top_ks) + + +def test_guard_refuses_fp64_temperatures() -> None: + args = _case(op.CELLS[0].spec, 2) + args[2] = args[2].double() + with pytest.raises(AssertionError, match="fp32 sampling parameters"): + with validating(op): + op(*args) + + +def test_the_verify_calls_this_targets_op_pair() -> None: + """The verify calls this target's op pair (the sm_107 tp4 kernel), whose workspace sizes as the shared one.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_sampled_accept as kernels, + ) + + assert kernels._sampled_accept_op is op + assert kernels._ws_words_op is ws_words + assert int(torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept_ws_words(4, 0)) == int( + torch.ops.trtllm.minimax_m3_sampled_accept_ws_words(4) + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_prewait_parameter_loads_change_nothing(cell) -> None: + """dispatch.SAMPLED_ACCEPT_PREWAIT_PARAMS moves only where the sampling parameters load (before or after the PDL + wait): the same draws and acceptance outputs, bit for bit, call after call.""" + outs = [] + for prewait in (True, False): + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + got = [] + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept(*args, prewait) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + outs.append(got) + for call, (a, b) in enumerate(zip(*outs)): + for x, y in zip(a, b): + assert torch.equal(x, y), f"call {call}: prewait changes {x.tolist()} -> {y.tolist()}" + + +def _walk_outputs(args, one_round_k, ilp, calls=3, offset_step=1000): + """The draws and acceptance outputs of `calls` calls with the owner walk one token (ilp False) or 8 tokens at a + time (ilp True), at SA-1R variant `one_round_k`, each call at its own Philox offsets.""" + if one_round_k: + words = torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept_ws_words(1, one_round_k) + args[6] = torch.zeros(words, dtype=torch.int64, device="cuda") + args[7] = args[6].data_ptr() + rows = args[0].shape[0] + got = [] + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * offset_step) + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept( + *args, True, False, one_round_k, None, ilp + ) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + return got + + +@pytest.mark.parametrize("one_round_k", [0, 82]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_ilp_walk_changes_nothing(cell, one_round_k) -> None: + """The owner walk's 8-token form (``ilp_walk``; off in this target) only computes the owner's token weights 8 at a + time; the walk adds them in the same order with the same early exit, so the draws and acceptance outputs are the + one-token walk's, bit for bit, call after call, with SA-1R off and on.""" + one = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, False) + ilp = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, True) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.95]) +def test_the_ilp_walk_changes_nothing_over_many_rounds(top_p) -> None: + """A flat tail (many tokens of equal weight) under a tiny and a wide nucleus: rows take many exchange rounds and + each round's walk exits at a different token; the 8-token walk draws the same tokens as the one-token walk.""" + spec = dict(nreq=1, vs=50016, vocab_offset=50016, check="flat", top_p=top_p, temperature=1.0) + one = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), 0, False, calls=24) + ilp = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), 0, True, calls=24) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +def test_the_ilp_walk_is_the_same_under_graph_replay() -> None: + """The 8-token walk captured in a CUDA graph and replayed draws what the one-token walk draws eagerly.""" + cell = op.CELLS[0] + eager = _walk_outputs(_case(cell.spec, seed_of(cell.why)), 0, False) + _walk_outputs( + _case(cell.spec, seed_of("ilp walk warm-up")), 0, True, calls=1 + ) # first launch outside the capture + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept(*args, True, False, 0, None, True) + torch.cuda.current_stream().wait_stream(stream) + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + graph.replay() + torch.cuda.synchronize() + for x, y in zip(eager[c], [args[i] for i in (9, 14, 15, 16, 18)]): + assert torch.equal(x, y), ( + f"replay {c}: the graphed 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_early_trigger_changes_nothing(cell) -> None: + """dispatch.SAMPLED_ACCEPT_EARLY_TRIGGER moves only where the grid releases its dependents: the same draws and + acceptance outputs, bit for bit, call after call, PDL launch on.""" + outs = [] + for early in (True, False): + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + got = [] + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept(*args, True, early) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + outs.append(got) + for call, (a, b) in enumerate(zip(*outs)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the early trigger changes {x.tolist()} -> {y.tolist()}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept_1r.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept_1r.py new file mode 100644 index 000000000000..42bacb1467be --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept_1r.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SA-1R (dispatch.SAMPLED_ACCEPT_1R_K) on trtllm::minimax_m3_sm107tp4_sampled_accept, world 1, the TP4 shard. + +The one-round variants (one_round_k 9: exact top-9 list; 82: pooled top-8 + bound) must return the two-round +kernel's (one_round_k 0) sampled tokens, accepted tokens, accepted counts, gather ids and extra-attempt count BIT FOR +BIT on the same Philox stream, over random AgentX-like rows (several temperatures / top-p, broad and heavy-tailed +heads) and over rows built so the decision sits within ulps of top_p * Z (the margin fallback) or a slice holds more +heavy tokens than its list (the incomplete-list fallback), eager and under CUDA-graph replay. The round-0 decided +counter must engage on the AgentX-like rows and stay below the row count on the boundary rows (fallbacks fire). +Ported from the prototype test (exact on 39,968 rows). +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + +VS, D = 50016, 3 +KX = (9, 82) + +if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the sm_107 tp4 sampled accept runs on SM 10.7", allow_module_level=True) + + +class Runner: + """One variant with its own workspace / counters (call tag sequence) and outputs.""" + + def __init__(self, k, rows): + self.k, self.rows = k, rows + words = int(torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept_ws_words(1, k)) + self.ws = torch.zeros(words, dtype=torch.int64, device="cuda") + self.counters = torch.zeros(2, dtype=torch.int32, device="cuda") + self.tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + self.extra = torch.zeros(1, dtype=torch.int32, device="cuda") + self.decided = torch.zeros(1, dtype=torch.int32, device="cuda") + nreq = rows // (D + 1) + self.acc = torch.zeros((nreq, D + 1), dtype=torch.int32, device="cuda") + self.nacc = torch.zeros(nreq, dtype=torch.int32, device="cuda") + self.gids = torch.zeros(nreq, dtype=torch.int64, device="cuda") + + def __call__(self, lg, temps, tps, seeds, offs, draft, bidx): + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept( + lg, 0, temps, tps, seeds, offs, self.ws, self.ws.data_ptr(), self.counters, self.tokens, 0, 1, draft, + bidx, self.acc, self.nacc, self.gids, 0, None, None, self.extra, False, True, False, self.k, + self.decided) # fmt: skip + + def out(self): + return (self.tokens[: self.rows].clone(), self.acc.clone(), self.nacc.clone(), self.gids.clone(), + self.extra.clone()) # fmt: skip + + +def _params(rows, temp, top_p, seed0): + temps = torch.full((rows,), temp, dtype=torch.float32, device="cuda") + tps = ( + (top_p if torch.is_tensor(top_p) else torch.full((rows,), top_p)) + .float() + .cuda() + .contiguous() + ) + seeds = torch.arange(rows, dtype=torch.int64, device="cuda") * 7919 + seed0 + offs = torch.zeros(rows, dtype=torch.int64, device="cuda") + return temps, tps, seeds, offs + + +def _sets(n, rows, g, kind): + out = [] + for _ in range(n): + lg = torch.randn(rows, VS, device="cuda", generator=g) * (1.0 if kind == "flat" else 2.0) + if kind == "lm": # a few heavy tokens over a broad tail + hot = torch.randint(0, VS, (rows, 8), device="cuda", generator=g) + lg.scatter_(1, hot, 12.0 + 3.0 * torch.rand(rows, 8, device="cuda", generator=g)) + elif kind == "tail": # 64 near-equal top tokens: lists overflow + hot = torch.randint(0, VS, (rows, 64), device="cuda", generator=g) + lg.scatter_(1, hot, 9.0 + 0.5 * torch.rand(rows, 64, device="cuda", generator=g)) + out.append(lg.to(torch.bfloat16).contiguous()) + return out + + +def _draft_from(lg, g): + nreq = lg.shape[0] // (D + 1) + am = lg.float().argmax(1).view(nreq, D + 1)[:, :D].int() + rnd = torch.randint(0, VS, (nreq, D), device="cuda", generator=g, dtype=torch.int32) + return torch.where(torch.rand(nreq, D, device="cuda", generator=g) < 0.7, am, rnd).contiguous() + + +def _compare(ref, got, what): + for n, x, y in zip( + ("tokens", "acc_tokens", "num_acc", "gather_ids", "extra_attempts"), ref, got + ): + assert torch.equal(x, y), f"{what}: {n} differs" + + +CASES = [(1.0, 0.95, "lm"), (0.7, 0.8, "lm"), (1.3, 0.5, "lm"), (1.0, 1.0, "lm"), (1.0, 0.95, "flat"), + (1.0, 0.95, "tail"), (0.6, 0.99, "tail")] # fmt: skip + + +def test_one_round_is_two_round_bitwise(): + g = torch.Generator(device="cuda").manual_seed(2) + total = 0 + decided = {k: 0 for k in KX} + for rows, calls in ((32, 20), (4, 60)): + for ci, (t, p, kind) in enumerate(CASES): + runners = {k: Runner(k, rows) for k in (0, *KX)} + temps, tps, seeds, offs = _params(rows, t, p, 1000 + ci) + sets = _sets(4, rows, g, kind) + bidx = torch.arange(rows // (D + 1), dtype=torch.int32, device="cuda") + for i in range(calls): + lg = sets[i % 4] + draft = _draft_from(lg, g) + offs.fill_(i) + for r in runners.values(): + r(lg, temps, tps, seeds, offs, draft, bidx) + for k in KX: + _compare( + runners[0].out(), runners[k].out(), f"K={k} rows={rows} case={ci} call={i}" + ) + total += rows + if kind == "lm" and p < 1.0: + for k in KX: + decided[k] += int(runners[k].decided.item()) + assert int(runners[0].decided.item()) == 0 + assert total >= 5000 + for k in KX: + assert decided[k] > 0, ( + f"K={k}: no row decided in round 0 on the AgentX-like rows (SA-1R not engaged)" + ) + + +def _boundary_rows(rows, n_heavy_same_slice): + """Heavy tokens heavier than a 50-token candidate group; top_p swept across ulps of S / Z (margin fallback), or + n heavy tokens in one slice (more than its list: incomplete fallback).""" + lg = torch.full((rows, VS), -40.0, device="cuda") + sl = ((VS + 7) // 8 + 7) // 8 * 8 + heavy = list(range(5, 5 + n_heavy_same_slice)) if n_heavy_same_slice else [5, 3 * sl + 11] + grp = [2 * sl + 3 + 37 * j for j in range(25)] + [5 * sl + 7 + 41 * j for j in range(25)] + lg[:, heavy] = 5.0 + lg[:, grp] = 4.0 + lg = lg.to(torch.bfloat16).contiguous() + w = torch.exp(lg[0].double() - lg[0].double().max()) + crit = float(w[heavy].sum()) / float(w.sum()) + ks = torch.arange(rows, dtype=torch.float64) - rows // 2 + return lg, (crit * (1.0 + ks * 2.0**-24)).float() + + +@pytest.mark.parametrize("graph", [False, True]) +def test_one_round_boundary_and_incomplete(graph): + rows = 32 + g = torch.Generator(device="cuda").manual_seed(3) + seen = {k: [0, 0] for k in KX} # decided rows, rows + for n_same in (0, 12): + lg, tps_rows = _boundary_rows(rows, n_same) + for shift in range(-4, 4): + tps = (tps_rows.double() * (1.0 + shift * 32 * 2.0**-24)).float() + runners = {k: Runner(k, rows) for k in (0, *KX)} + temps, tp_t, seeds, offs = _params(rows, 1.0, tps, 5000 + shift) + bidx = torch.arange(rows // (D + 1), dtype=torch.int32, device="cuda") + draft = _draft_from(lg, g) + graphs = {} + if graph: + for r in runners.values(): + r(lg, temps, tp_t, seeds, offs, draft, bidx) + torch.cuda.synchronize() + for k, r in runners.items(): + graphs[k] = torch.cuda.CUDAGraph() + with torch.cuda.graph(graphs[k]): + r(lg, temps, tp_t, seeds, offs, draft, bidx) + for r in runners.values(): + r.decided.zero_() + for i in range(12): + offs.fill_(i) + for k, r in runners.items(): + if graph: + graphs[k].replay() + else: + r(lg, temps, tp_t, seeds, offs, draft, bidx) + for k in KX: + _compare(runners[0].out(), runners[k].out(), + f"boundary n_same={n_same} shift={shift} call={i} K={k}") # fmt: skip + for k in KX: + seen[k][0] += int(runners[k].decided.item()) + seen[k][1] += 12 * rows + for k in KX: + assert 0 < seen[k][0] < seen[k][1], ( + f"K={k}: decided {seen[k]}: the boundary sweep missed both sides" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..0c90ceac2e70 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_shard_accept entry (greedy verify from this rank's logits shard).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + craft_drafts, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" + + +def _case(spec, seed, edit_logits=None): + g = gen(seed) + nreq, d, vs, world, rank = ( + spec["nreq"], + spec["draft_len"], + spec["vs"], + spec["world"], + spec["rank"], + ) + m = nreq * (d + 1) + if spec["strided"]: + logits = randn((m, vs + 72), g, 3.0)[:, :vs] + else: + logits = randn((m, vs), g, 3.0) + if edit_logits is not None: + edit_logits(logits) + uc, keys, counters = head_workspace(world, slot_rows=SHARD_ROWS) + own = logits.float().max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=SHARD_ROWS, ties=spec.get("ties", False) + ) + offset = rank * vs + best, idx = local_argmax(logits.float(), offset) + tokens = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, SHARD_ROWS) + draft = craft_drafts(tokens, nreq, d) + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + force_total = 3 if spec["force"] == "total" else 0 + force_vec = (torch.randint(1, d + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if spec["force"] == "vec" else None) # fmt: skip + return (logits, offset, keys, counters, uc, uc.data_ptr(), rank, world, draft, batch_idx, acc, num, gather, + force_total, next_new, num_sms(), force_vec) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bit-exact: the op reduces BF16 logits without arithmetic, ties (within the shard and across ranks) go to + the lowest global index on both sides.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) + assert args[3].tolist() == [0, 1] and not bool(args[2].any()) + + +def test_guard_refuses_nan_logits() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0][1, 7] = float("nan") + with pytest.raises(AssertionError, match="NaN"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_batch_index_past_the_requests() -> None: + args = list(_case(op.CELLS[0].spec, 2)) + args[9] = args[9] + 1 + with pytest.raises(AssertionError, match="batch_idx"): + with validating(op): + op(*args) + + +def _signed_zero_maxima(logits: torch.Tensor) -> None: + """Every row's maximum is zero, reached as -0.0 at a lower index and +0.0 at a higher one.""" + logits.fill_(-1.0) + logits[:, 3] = -0.0 + logits[:, 9] = 0.0 + + +@pytest.mark.parametrize("ties", [False, True], ids=["peer_rows", "peer_ties"]) +def test_signed_zero_maxima_take_the_lower_index(ties) -> None: + """-0.0 and +0.0 compare equal, so the lower index wins as with torch.argmax (also across ranks).""" + spec = dict(op.CELLS[1].spec, ties=ties) + args = _case(spec, seed_of(f"signed-zero-{ties}"), edit_logits=_signed_zero_maxima) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..4b27e42af50b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_slot_scatter entry (the spec sampler's per-slot stores in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_slot_scatter import ( # noqa: E501 + m3_slot_scatter as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" + +STALE = -4 + + +def _case(spec, seed): + """Worker outputs and the store's transposed [slots, width] views of its [width, slots, 1] buffers.""" + g = gen(seed) + n, nslots = spec["n"], spec["slots"] + + def rnd(shape): + return torch.randint(0, VOCAB, shape, generator=g, device="cuda", dtype=torch.int32) + + store_new = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_next = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_lens = torch.full((nslots,), STALE, dtype=torch.int32, device="cuda") + store_draft = torch.full((nslots, spec["store_d"]), STALE, dtype=torch.int32, device="cuda") + slots = torch.randperm(nslots, generator=g, device="cuda")[:n] + lens = torch.randint(1, 5, (n,), generator=g, device="cuda", dtype=torch.int32) + return (slots, rnd((n, spec["src_w"])), store_new.squeeze(-1).T, rnd((n, spec["src_w"])), + store_next.squeeze(-1).T, lens, store_lens, rnd((n, spec["src_d"])), store_draft) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, zero padding included, other slots untouched.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in (("da", args[2]), ("db", args[4]), ("dl", args[6]), ("dd", args[8])): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_repeated_slots() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0] = args[0].clone() + args[0][1] = args[0][0] + with pytest.raises(AssertionError, match="distinct"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py new file mode 100644 index 000000000000..4a2a242816d3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for m3_spec_offsets (RuntimeHooks.spec_offsets_op): the overlap scheduler's two speculative step corrections +in one plain launch, bit for bit the in-place ATen adds, eager and in a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_spec_offsets import ( # noqa: E501 + m3_spec_offsets as op, +) + +assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" + +I32 = torch.iinfo(torch.int32) + + +def _inputs(tokens, seqs, seed, wrap=False): + g = torch.Generator(device="cuda").manual_seed(seed) + lo, hi = (I32.min, I32.max) if wrap else (0, 1 << 20) + pos = torch.randint(lo, hi, (tokens,), generator=g, device="cuda", dtype=torch.int32) + kv = torch.randint(lo, hi, (seqs,), generator=g, device="cuda", dtype=torch.int32) + pos_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (tokens,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + kv_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (seqs,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + if wrap: + pos[:2], pos_off[:2] = I32.max, 1 # overflow wraps, as the ATen add does + kv[:1], kv_off[:1] = I32.min, -1 + return pos, pos_off, kv, kv_off + + +CELLS = [c.spec for c in op.CELLS] + [ + dict(tokens=1, seqs=1), + dict(tokens=300, seqs=75), + dict(tokens=4, seqs=0), +] + + +@pytest.mark.parametrize("spec", CELLS, ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items())) +def test_bitwise_aten_adds(spec) -> None: + for trial in range(3): + pos, pos_off, kv, kv_off = _inputs( + spec["tokens"], spec["seqs"], 97 * trial + spec["tokens"], spec.get("wrap", False) + ) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + assert op(pos, pos_off, kv, kv_off) is True + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"trial {trial}" + + +def test_graph_replays_read_the_live_offsets() -> None: + pos, pos_off, kv, kv_off = _inputs(32, 8, 5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + for step in range(3): + fresh = _inputs(32, 8, 11 + step) + for dst, src in zip((pos, pos_off, kv, kv_off), fresh): + dst.copy_(src) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"step {step}" + + +def test_runner_views_are_accepted() -> None: + """The runner passes position_ids[0, ctx:] and kv_lens_cuda[ctx_reqs:seqs]: 1-D slices of contiguous buffers.""" + position_ids = torch.arange(40, dtype=torch.int32, device="cuda").view(1, 40) + kv_lens = torch.arange(12, dtype=torch.int32, device="cuda") + pos_off = torch.full((40,), 2, dtype=torch.int32, device="cuda") + kv_off = torch.full((12,), 3, dtype=torch.int32, device="cuda") + assert op(position_ids[0, 8:], pos_off[:32], kv_lens[2:10], kv_off[:8]) is True + torch.cuda.synchronize() + assert torch.equal( + position_ids[0], torch.cat([torch.arange(8), torch.arange(8, 40) + 2]).int().cuda() + ) + assert torch.equal( + kv_lens, + torch.cat([torch.arange(2), torch.arange(2, 10) + 3, torch.arange(10, 12)]).int().cuda(), + ) + + +@pytest.mark.parametrize( + "bad", + ["int64", "strided", "length", "2d"], +) +def test_declines_other_layouts_without_writing(bad) -> None: + pos, pos_off, kv, kv_off = _inputs(8, 2, 3) + if bad == "int64": + kv = kv.long() + elif bad == "strided": + pos = torch.zeros(16, dtype=torch.int32, device="cuda")[::2] + elif bad == "length": + pos_off = pos_off[:7] + else: + pos = pos.view(2, 4) + before = (pos.clone(), kv.clone()) + assert op(pos, pos_off, kv, kv_off) is False + torch.cuda.synchronize() + assert torch.equal(pos, before[0]) and torch.equal(kv, before[1]) + + +def test_launches_plainly() -> None: + """It writes step inputs that pre-wait readers take before their grid wait (MSA contract, dispatch + PRODUCER_PREWAIT_STEP_INPUTS): its kernel node carries no programmatic (PDL) serialization.""" + from cuda.bindings import runtime as rt + + pos, pos_off, kv, kv_off = _inputs(32, 8, 7) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + err, _, n = rt.cudaGraphGetNodes(raw, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + flags = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind == rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + flags.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + assert flags == [0], flags diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..c57435ce1f30 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_gather entry (the overlap scheduler's step-input gathers in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_step_gather import ( # noqa: E501 + m3_step_gather as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" + +#: Rows of the engine's flat buffers (the max batch), and their stale fill. +MAX_ROWS = 8 +STALE = -4 + + +def _case(spec, seed): + g = gen(seed) + p, t, d, nslots, b0 = spec["p"], spec["t"], spec["d"], spec["slots"], spec["b0"] + i32 = dict(dtype=torch.int32, device="cuda") + new_tok = torch.randint(0, VOCAB, (t, nslots, 1), generator=g, device="cuda", dtype=torch.int32) + if spec.get("two_d"): + new_tok = new_tok.squeeze(-1) + next_draft = torch.randint(0, VOCAB, (nslots, d), generator=g, device="cuda", dtype=torch.int32) + lens = torch.randint(1, t + 1, (nslots,), generator=g, device="cuda", dtype=torch.int32) + slots = torch.randperm(nslots, generator=g, device="cuda")[:p].to(torch.int32) + ppi = slots.repeat_interleave(t) + ids = torch.full((MAX_ROWS * t,), STALE, **i32) + draft = torch.full((MAX_ROWS * d,), STALE, **i32) + pos_off = torch.full((MAX_ROWS * t,), STALE, **i32) + kv_off = torch.full((MAX_ROWS,), STALE, **i32) + return ( + new_tok, + next_draft, + lens, + slots, + ppi, + p, + t, + d, + ids, + spec["ids0"], + draft, + spec["d0"], + pos_off, + kv_off, + b0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in ( + ("ids", args[8]), + ("draft", args[10]), + ("pos_off", args[12]), + ("kv_off", args[13]), + ): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_a_slot_past_the_store() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[3] = args[3].clone() + args[3][0] = 8 + with pytest.raises(AssertionError, match="slots must index the store"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py new file mode 100644 index 000000000000..a5d0cec0ed6c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_prologue entry (the decode graph's step-input prologue in one launch).""" + +import random +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_step_prologue import ( # noqa: E501 + m3_step_prologue as op, +) +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" + +MAX_DESC = StepInputPack.MAX_DESC +WINDOW = StepInputPack.GRAPH_WINDOW +DATA_START = StepInputPack.HEADER_BYTES + StepInputPack.DESC_BYTES * MAX_DESC +ARENA = 1 << 20 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _layout(spec, seed): + """Copies as (dst, src, bytes) arena offsets (a host copy's src: a window offset) and the host copies' bytes.""" + rng = random.Random(seed) + h2d, d2d, data = [], [], [] + used = DATA_START + pos = 0 + for _ in range(spec["h2d"]): + n = rng.choice([1, 3, 4, 7, 16, 24, 32, 48, 64, 5, 12]) + pos += rng.randint(0, 40) # gaps and every alignment + h2d.append((pos, used, n)) + data.append(bytes(rng.getrandbits(8) for _ in range(n))) + pos += n + used += -(-n // 16) * 16 + assert pos < ARENA // 2 and used <= WINDOW + # Device copies write the last quarter; unchained ones read the third. + dpos, spos = 3 * ARENA // 4, ARENA // 2 + for i in range(spec["d2d"]): + if spec.get("chained") and h2d: + dst_h, _, n_h = h2d[i % len(h2d)] + off = rng.randint(0, n_h - 1) + src, n = dst_h + off, n_h - off + else: + n = rng.choice([4, 8, 16, 32, 6, 64]) + spos += rng.randint(0, 24) + src = spos + spos += n + dpos += rng.randint(0, 24) + d2d.append((dpos, src, n)) + dpos += n + assert dpos < ARENA and spos < 3 * ARENA // 4 + return h2d, d2d, data + + +def _window(layout, base: int) -> torch.Tensor: + """The window of ``layout`` for an arena at device address ``base``.""" + h2d, d2d, data = layout + w = torch.zeros(WINDOW, dtype=torch.uint8) + words = [len(h2d), len(d2d)] + for dst, src, n in h2d: + words += [base + dst, src, n] + for dst, src, n in d2d: + words += [base + dst, base + src, n] + w[: 8 * len(words)] = torch.tensor(words, dtype=torch.int64).view(torch.uint8) + for (_, off, n), b in zip(h2d, data): + w[off : off + n] = torch.frombuffer(bytearray(b), dtype=torch.uint8) + return w.cuda() + + +def _arenas(seed): + a = torch.randint(0, 256, (ARENA,), generator=gen(seed), device="cuda", dtype=torch.int32) + a = a.to(torch.uint8) + return a, a.clone() + + +@pytest.mark.parametrize("with_epoch", [True, False], ids=["epoch", "no_epoch"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, with_epoch, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Byte for byte the default prologue's destinations, header clear and epoch.""" + seed = seed_of(cell.why) + layout = _layout(cell.spec, seed) + arena, ref_arena = _arenas(seed) + window, ref_window = _window(layout, arena.data_ptr()), _window(layout, ref_arena.data_ptr()) + epoch = torch.full((1,), 41, dtype=torch.int64, device="cuda") if with_epoch else None + ref_epoch = epoch.clone() if with_epoch else None + certifying(op) + op(window, epoch, MAX_DESC) + op.reference(ref_window, ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena) + # Each window holds its own arena's addresses: past the consumed header, both are left as staged. + for w, a in ((window, arena), (ref_window, ref_arena)): + assert not w[:16].any(), "the header is consumed" + assert torch.equal(w[16:], _window(layout, a.data_ptr())[16:]) + if with_epoch: + assert int(epoch.item()) == int(ref_epoch.item()) == 42 + + +def test_graph_replays_scatter_each_new_window() -> None: + """Captured once on a static window: every replay writes that replay's window and bumps the epoch once.""" + arena, ref_arena = _arenas(7) + window = torch.zeros(WINDOW, dtype=torch.uint8, device="cuda") + epoch = torch.zeros(1, dtype=torch.int64, device="cuda") + ref_epoch = epoch.clone() + op(window, epoch, MAX_DESC) # the empty window: writes nothing + op.reference(window.clone(), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(window, epoch, MAX_DESC) + for step, spec in enumerate( + [dict(h2d=24, d2d=8, chained=True), dict(h2d=0, d2d=0), dict(h2d=40, d2d=4)] + ): + layout = _layout(spec, 100 + step) + window.copy_(_window(layout, arena.data_ptr())) + graph.replay() + op.reference(_window(layout, ref_arena.data_ptr()), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena), step + assert not window[:16].any(), step + assert int(epoch.item()) == int(ref_epoch.item()) == step + 2, step + + +def _pack_case(prologue_op): + pack = StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=True, + in_graph=True, + early_copy=True, + step_epoch=True, + prologue_op=prologue_op, + ) + a, b, c = (torch.zeros(16, dtype=torch.int32, device="cuda") for _ in range(3)) + out = torch.zeros(48, dtype=torch.int32, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _pack_step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + + def host(v): + return torch.tensor(v, dtype=torch.int32).pin_memory() + + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, host(values[0])) + step_inputs.copy_(b, host(values[1])) + if overlap: + step_inputs.copy_(a[:4], host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if not flush: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_pack_with_the_one_launch_prologue_matches_the_default() -> None: + """The same pack steps through the default prologue and the one-launch one: the same outputs and epochs, and the + one-launch prologue is engaged (the default's warm-up compile is not what bumped its epoch).""" + packs = [_pack_case(None), _pack_case(op)] + assert packs[0][0].prologue_op is None and packs[1][0].prologue_op is op + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + outs = [] + for pack, graph, dsts, out in packs: + # A replay with a new pack, one with an eagerly issued pack, and one with nothing to scatter. + _pack_step(pack, graph, dsts, (va, vb), overlap=step == 2, flush=step != 3) + torch.cuda.synchronize() + outs.append((out.tolist(), int(pack.epoch.item()))) + assert outs[0] == outs[1], step + assert outs[1][1] == step + 2, step # one bump at warm-up, then one per replay diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..d8eb4f29e589 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The feed-forward exit on Eagle3 capture layers: the generic layer in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``), and the decode fast path's capture layers (``CAPTURE_LAYERS_FAST``). + +Every layer's MoE / MLP output goes unreduced to the fused boundary (all-reduce + residual + the next layer's input +norm, the final norm after the last layer), and a capture layer captures the reduced residual sum that boundary +writes, with no residual of its own. CPU only: the attention, the feed-forward and the boundaries are stand-ins. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) + +LAYERS = 4 + + +class _Spec: + def __init__(self, layers): + self.layers = set(layers) + + def is_layer_capture(self, i): + return i in self.layers + + +@pytest.fixture(params=["prefill", "decode"]) +def layer(request, monkeypatch): + """(core, the target's generic layer bound to a target over ``core``).""" + calls = [] + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + c = SimpleNamespace( + w=w, num_layers=LAYERS, moe_layers=frozenset({2, 3}), rms_eps=1e-6, calls=calls + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, + "_native_allreduce_norm", + lambda core, x, r, norm: (("normed", x, norm), ("updated", x, r)), + ) + + def boundary(self, x, x2, residual, norm_w): + calls.append((x, x2, norm_w)) + return ("next", norm_w), ("sum", x, x2) + + if request.param == "prefill": + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i))) + monkeypatch.setattr( + cls, "_swiglu_mlp", lambda self, h, prefix, low_m, dense=False: ("mlp", prefix) + ) + else: + cls = modeling.DecodeTarget + monkeypatch.setattr( + cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i), False, False) + ) + monkeypatch.setattr( + cls, + "_swiglu_mlp", + lambda self, h, prefix, low_m, push=None, dense=False: (("mlp", prefix), False), + ) + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + monkeypatch.setattr(cls, "_generic_boundary", boundary) + target = cls(c) + return c, target._layer if request.param == "prefill" else target._generic_layer + + +@pytest.mark.parametrize("i", range(LAYERS)) +def test_a_capture_layer_captures_the_boundary_residual_sum(layer, i) -> None: + core, run = layer + captured = [] + out = run( + i, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([i]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + norm = core.w[f"l{i + 1}_in_norm"] if i + 1 < LAYERS else "final" + assert core.calls[-1][2] == norm, ( + "the boundary applies the next layer's norm (the final norm after the last)" + ) + assert captured == [(i, out[1], None)] + if i in core.moe_layers: + assert core.calls[-1][1] == ("routed", i), "the routed part folds into the boundary" + + +def test_a_layer_without_capture_captures_nothing(layer) -> None: + _, run = layer + captured = [] + run( + 1, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([]), + lambda *a: captured.append(a), + ) + assert not captured + + +class _FastCore: + num_layers = 3 + w = {"embed": None} + + +@pytest.mark.parametrize("written", [True, False]) +def test_a_decode_capture_layer_takes_the_fast_layer_with_the_capture_norm( + monkeypatch, written +) -> None: + """A decode step's capture layer (past the first) runs the fast decode layer with the worker's capture target, and + captures the boundary's updated residual; ``written=True`` when the boundary all-reduce wrote the capture norm + (``CAPTURE_NORM_IN_BOUNDARY``).""" + assert modeling.CAPTURE_LAYERS_FAST and modeling.CAPTURE_FUSED_BOUNDARY + assert modeling.CAPTURE_NORM_IN_BOUNDARY + target = modeling.DecodeTarget(_FastCore()) + seen, captured = [], [] + monkeypatch.setattr(target, "_check_step_contract", lambda md: None) + monkeypatch.setattr(target, "step_args", lambda md: {"num_tokens": 4}) + monkeypatch.setattr(target, "_embed", lambda ids, n: ("h", "r")) + monkeypatch.setattr(modeling, "plan_for", lambda n: SimpleNamespace(name="plan")) + monkeypatch.setattr(modeling, "_record", lambda *a: None) + + def fast(i, plan, pos, h, r, step, capture=None): + seen.append((i, capture)) + if capture is not None: + target._capture_written = written + return ("h", i), ("r", i) + + def generic(*a): + raise AssertionError("a capture layer ran the generic layer on a decode step") + + monkeypatch.setattr(target, "_layer", fast) + monkeypatch.setattr(target, "_generic_layer", generic) + out = target.forward( + None, + torch.zeros(4, dtype=torch.int32), + None, + None, + _Spec([1]), + lambda spec, i, h, r, written=False: captured.append((i, h, r, written)), + lambda spec, i, n: ("capture_out", i, n), + ) + assert out == ("h", 2) + assert seen == [(0, None), (1, ("capture_out", 1, 4)), (2, None)] + assert captured == [(1, ("r", 1), None, written)] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..7b0535e70812 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,286 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp4 (Rubin) decode target's fast-path table (``modeling.py``: the row ranges and the selection functions): its +row ranges against the catalog cells that certify them, the selection for every decode row count, and the phase split. + +The paths whose cells pass on SM 10.7 take the sm_103 tp4 rows; the ones built on kernels with no SM 10.7 build (the +the fp32 routers, the low-M GEMM, the fused / mega / c1 attention) and the ones not yet driven on +SM 10.7 (the long-row select, the capped dense decode split, the early drafter fc) cover nothing. The dense layers' and +the drafter's full-context decode take the one-request shapes measured at TP4. + +CPU only: nothing here launches a kernel. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + triton_sparse_decode as tsd, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm import ( # noqa: E501 + m3_mx_gemv, +) + +ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 2) # every decode row count and the first one past it + +# The TP4 per-rank MXFP8 projections (N, K): 16 query heads and 1 KV head, expert width 768 per rank. +TP4_PROJECTIONS = ( + m3_mx_gemv.QKV_SPARSE, + m3_mx_gemv.QKV, + m3_mx_gemv.O_PROJ, + m3_mx_gemv.SHARED_GATE_UP, + m3_mx_gemv.SHARED_DOWN, + m3_mx_gemv.DENSE_GATE_UP, + m3_mx_gemv.DENSE_DOWN, +) + + +@pytest.fixture(autouse=True) +def _policy(): + """Every test starts outside require mode with nothing logged, and leaves it so.""" + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_every_covered_shape_is_a_certified_cell(fast_path) -> None: + uncertified = fast_path.keys - fast_path.certified() + assert not uncertified, ( + f"{fast_path.path} covers {sorted(uncertified, key=str)} but " + f"{type(fast_path.entry).__name__}.CELLS certifies none of them: add a cell or narrow the range" + ) + + +def test_the_fast_path_names_are_unique() -> None: + names = [r.path for r in dispatch.FAST_PATHS] + assert len(names) == len(set(names)) + + +def test_the_decode_target_keeps_the_tp4_geometry() -> None: + """Up to 8 requests verifying 1 + 3 drafts and their padding rows; a wider step is the prefill target's.""" + assert dispatch.DECODE_MAX_TOKENS == 48 + assert [dispatch.is_decode_step(n) for n in ROWS] == [True] * 48 + [False] + + +# ------------------------------------------------------------------ the paths still off on SM 10.7 + +OFF_ON_SM107 = ( + dispatch.DENSE_DECODE_SPLIT, + dispatch.LONG_SELECT, + dispatch.MC_FINALIZE, +) + + +@pytest.mark.parametrize("fast_path", OFF_ON_SM107, ids=lambda r: r.path) +def test_paths_without_an_sm107_build_cover_nothing(fast_path) -> None: + assert not fast_path.keys, f"{fast_path.path} covers {sorted(fast_path.keys, key=str)}" + + +def test_decode_boundaries_follow_the_multicast_all_reduce_kill() -> None: + """MC_AR_TP4 moves every decode boundary (and the o_proj / MoE pushes into it) onto the SM 10.7 TP4 multicast + all-reduce; off, the built-in all-reduce runs and nothing pushes.""" + on = dispatch.MC_AR_TP4 + for plan in dispatch.DECODE_PLANS: + assert (plan.attn_mc_ar, plan.moe_mc_ar, plan.qkv_prequant) == (on, on, on), plan + assert tuple(p.max_tokens for p in dispatch.DECODE_PLANS) == (dispatch.DECODE_MAX_TOKENS,) + assert bool(dispatch.MC_ALLREDUCE.keys) == on + assert dispatch.DRAFT_MC_AR.keys == (frozenset(range(1, 9)) if on else frozenset()) + assert dispatch.ATTN_PUSH_MAX_TOKENS == (8 if on else 0) + assert dispatch.DENSE_PUSH_MAX_TOKENS == (8 if on and dispatch.DENSE_PUSH_TP4 else 0) + assert dispatch.MOE_PUSH_MAX_TOKENS == (dispatch.MOE_ENGINE_MAX_TOKENS if on else 0) + assert dispatch.GENERIC_MC_AR_MAX_TOKENS == 0 + assert dispatch.GENERIC_MC_AR_FOLDED_MAX_TOKENS == 0 + + +def test_the_dense_decode_split_is_the_built_in_one() -> None: + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {} + + +# ------------------------------------------------------------------ on: the dense layers' full-context decode + + +def test_the_dense_layers_take_the_fullctx_decode_at_the_measured_tp4_shapes() -> None: + """One or two requests, 1 + 3 verify rows or one draft-pass row, each at its own pinned recipe.""" + assert dispatch.FULLCTX_DECODE_ENABLED + assert dispatch.FULLCTX_DECODE.keys == {(4, 1), (1, 1), (4, 2), (1, 2)} + assert set(dispatch.FULLCTX_DECODE_RECIPES) == set(dispatch.FULLCTX_DECODE.keys) + for recipe in dispatch.FULLCTX_DECODE_RECIPES.values(): + assert 2 <= recipe["splits"] <= 128 # the standalone reducer's range + + +def test_the_fullctx_plans_keep_the_tp4_geometry() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + + assert (kernel.HEADS, kernel.KV_HEADS, kernel.HEAD_DIM) == (16, 1, 128) + assert kernel.HEADS_PER_KV == 16 + + +def test_the_fullctx_entry_cells_run_at_the_forward_plan_capacity() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention import ( # noqa: E501 + m3_fullctx_decode as entry, + ) + + assert entry.MAX_KV_LEN == dispatch.FULLCTX_DECODE_MAX_KV_LEN + + +def test_the_fullctx_capacity_covers_the_agentx_cache_bound() -> None: + """The forward compares the cache's length bound (max_seq_len plus the speculative rows it reserves) with the plan + capacity; at the AgentX max_seq_len of 1 << 20 with up to 8 draft rows the dense layers must still be taken.""" + assert dispatch.FULLCTX_DECODE_MAX_KV_LEN >= (1 << 20) + 8 + + +def test_the_drafter_takes_the_fullctx_decode_at_the_same_shapes() -> None: + assert dispatch.DRAFTER_FULLCTX_ENABLED + assert dispatch.DRAFTER_FULLCTX.keys == dispatch.FULLCTX_DECODE.keys + + +def test_the_dense_fp8_qkv_and_kv_insert_kill_constants_default_on() -> None: + assert dispatch.DENSE_QKV_FP8_VIA_BF16 is True + assert dispatch.DENSE_KV_INSERT_FUSED is True + assert dispatch.FULLCTX_DECODE_OWN_REDUCER is True + assert dispatch.FULLCTX_DECODE_MXFP8 is True + + +def test_the_fullctx_taken_lines_log_without_raising_under_require() -> None: + dispatch.set_require(True) + dispatch.taken(dispatch.FULLCTX_DECODE.path, 4) + dispatch.taken(dispatch.DENSE_KV_INSERT_PATH, 4) + dispatch.taken(dispatch.FULLCTX_DECODE_MXFP8_PATH, 4) + + +# ------------------------------------------------------------------ the paths on + + +@pytest.mark.parametrize("rows", range(1, dispatch.DECODE_MAX_TOKENS + 1)) +def test_the_moe_takes_the_engines_and_the_generation_router(rows) -> None: + """1..8 rows the MoE engine, 9..13 the MoE16 engine (routers in-kernel); the unfused MoE of a pure generation step + takes the TF32 tensor-core router at 9..16 rows (14..16 always, 9..13 on an MoE16 miss), the fp32 router else.""" + expected = ( + dispatch.MOE_ENGINE.path + if rows <= 8 + else dispatch.MOE16_ENGINE.path + if rows <= 13 + else None + ) + assert dispatch.moe_path(rows) == expected + assert dispatch.ROUTER_TC_GEN.covers(rows) == (9 <= rows <= 16) + + +def _router_core(cluster16: bool): + import torch + + return SimpleNamespace(router_tc_cluster16=cluster16, w={"l3_router": torch.zeros(128, 64)}) + + +@pytest.mark.parametrize("rows", [1, 8, 9, 13, 14, 16, 17, 48]) +@pytest.mark.parametrize("cluster16", [True, False]) +def test_the_generation_router_takes_router_tc_only_on_its_rows( + monkeypatch, rows, cluster16 +) -> None: + """The built-in gate's choice: router_tc on 9..16 bf16 rows when the device co-schedules 16-CTA clusters; the fp32 + router otherwise. Context steps (``_router_logits``) never take it.""" + import torch + + calls = [] + monkeypatch.setattr(dispatch, "m3_router_tc", lambda *a: calls.append("tc") or [a[0]]) + monkeypatch.setattr(dispatch, "m3_fp32_linear", lambda *a: calls.append("fp32") or a[0]) + h = torch.zeros(rows, 64, dtype=torch.bfloat16) + dispatch._decode_router_logits(_router_core(cluster16), 3, h) + assert calls == ["tc" if cluster16 and 9 <= rows <= 16 else "fp32"] + calls.clear() + dispatch._router_logits(_router_core(cluster16), 3, h) + assert calls == ["fp32"] + + +def test_the_core_selects_the_target_on_the_context_count_alone() -> None: + """Every pure generation step is decode's, whatever its width; a step with context rows (a mixed one too) is + prefill's.""" + core = SimpleNamespace(_prefill="prefill", _decode="decode") + select = dispatch.MiniMaxM3Tp4ModelingV2._select_target + assert select(core, SimpleNamespace(num_contexts=0)) == "decode" + assert select(core, SimpleNamespace(num_contexts=1)) == "prefill" + assert select(core, SimpleNamespace(num_contexts=7)) == "prefill" + + +@pytest.mark.parametrize("rows", range(1, dispatch.DECODE_MAX_TOKENS + 1)) +def test_the_projections_take_the_gemvs_else_the_mxfp8_gemm(rows) -> None: + for nk in TP4_PROJECTIONS: + prequant = dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) + assert prequant == ( + dispatch.MX_GEMV_PATH + if rows in dispatch.MX_GEMV_ROWS[nk] + else dispatch.PREQUANT_GEMM_PATH + ) + bf16 = dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) + assert bf16 == ( + dispatch.MX_GEMV_BF16_PATH + if rows <= dispatch.MX_GEMV_BF16_MAX_TOKENS[nk] + else dispatch.GENERIC_GEMM_PATH + ) + + +def test_the_attention_producer_and_scorer_are_on() -> None: + assert dispatch.PRODUCER.covers("nvfp4_tp4") + assert dispatch.INDEX_SCORE.keys == {"early", "prefetch"} + assert dispatch.SPARSE_DECODE_PRODUCER.keys == {"per_row", "multi_query"} + assert dispatch.EMBED_FIRSTNORM.keys == frozenset(range(1, 9)) + + +def test_speculation_runs_the_fused_tail_with_the_bf16_draft_head() -> None: + requests = dispatch.DECODE_MAX_TOKENS // (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert dispatch.SHARD_ACCEPT_MAX_REQUESTS == requests + assert dispatch.SAMPLED_VERIFY_MAX_REQUESTS == requests + assert dispatch.FIRST_DRAFTER_INPUTS_MAX_REQUESTS == requests + assert dispatch.TARGET_HEAD_MAX_ROWS == 0 # the in-tree BF16 head is faster at the TP4 shard + assert dispatch.FC_AFTER_HEAD_MAX_ROWS == 8 + assert dispatch.DRAFT_STEP_MAX_REQUESTS == 8 + assert dispatch.FC_EARLY_MAX_ROWS == 0 + assert not dispatch.DRAFT_HEAD_CODED + + +def test_a_miss_raises_only_under_require() -> None: + dispatch.missed("some_path", 1, "a precondition") + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="some_path"): + dispatch.missed("some_path", 1, "a precondition") + + +def test_taken_logs_without_raising_under_require() -> None: + dispatch.set_require(True) + dispatch.taken(dispatch.INDEX_SCORE.path, 1) + dispatch.taken(dispatch.MOE_ENGINE.path, 8) + + +def test_taken_logs_once_per_path_and_rows(monkeypatch) -> None: + lines = [] + monkeypatch.setattr(dispatch.logger, "info", lines.append) + for _ in range(3): + dispatch.taken(dispatch.MX_GEMV_PATH, 4, "[2560, 6144] mxfp8 in") + dispatch.taken(dispatch.MX_GEMV_PATH, 8, "[2560, 6144] mxfp8 in") + assert len(lines) == 2 and all("m3_mx_gemv taken at" in line for line in lines) + + +def test_the_sparse_decode_writes_the_mxfp8_handoff() -> None: + """Every decode step (at most DECODE_MAX_TOKENS rows) can take the MXFP8 o_proj hand-off of the sparse decode.""" + assert dispatch.MERGE_MXFP8_HANDOFF is True + assert tsd.MERGE_MXFP8_MIN_TOKENS <= 1 and dispatch.DECODE_MAX_TOKENS <= tsd.MergeMxfp8.MAX_ROWS + + +def test_the_decode_mxfp8_decision_is_pinned_to_cute_dsl() -> None: + from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, + ) + + assert dispatch.MX_TAC_TP4 + assert dispatch.MXFP8_DECODE_QUANT_TACTIC == MXFP8QuantizeRunner.CUTE_DSL + assert dispatch.MXFP8_DECODE_GEMM_TACTIC == FlashInferMXFP8GemmRunner.CUTE_DSL diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py new file mode 100644 index 000000000000..2c76156860e8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode kernels the core builds before the executor's warmup (``prebuild_kernels``), and that a forward after it +never compiles one: the MoE engines of every ``MOE_ENGINE`` / ``MOE16_ENGINE`` token count, the decode GEMVs with +their push variants, the index scorer with and without its early-issue schedule at every decode query length (and +the T2SC scorer, the long-row select and the full-context plans where they apply), and the 9..16-row generation +router's device decision. + +The engine builds are stubbed (a real one needs a loaded checkpoint's expert buffers; the engine kernels have their +own tests); the GEMVs and the scorer compile for real on the GPU. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers the trtllm ops +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.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_107__tp4.kernels import ( # noqa: E501 + m3_index_score_t2sc as index_score_t2sc_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, + SHARED_DOWN, + SHARED_GATE_UP, +) + +dispatch = modeling + +assert torch.cuda.is_available(), "the prebuild checks require a CUDA device" + +INDEX_HEADS = 4 +INDEX_DIM = 128 +DRAFT_LEN = 3 +_CORE_METHODS = ( + "prebuild_kernels", + "_prebuild_engines", + "_try_build", + "engine", + "moe16_engine", + "_lookup", + "_scorer_variants", + "_prebuild_index_scorer", + "_prebuild_index_scorer_t2sc", + "decode_query_lens", +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require, frozen = dispatch.require_mode(), mx_gemv.frozen() + yield + dispatch.set_require(require) + mx_gemv.freeze(frozen) + dispatch.reset_logged() + + +class _FakeSparseAttn: + def __init__(self, indexer_kv_dtype: str): + self.indexer_kv_dtype = indexer_kv_dtype + self.m3_config = SimpleNamespace(num_index_heads=INDEX_HEADS, sparse_index_dim=INDEX_DIM) + + +def _core(num_layers: int = 3, indexer_kv_dtype: str = "fp8", built=None, fail=False): + """The parts of ``MiniMaxM3Tp4ModelingV2`` the prebuild reads: one rank (the build agreement takes no + collective), MoE + sparse layers 1.. (layer 0 dense), no projections (the GEMV prebuild has its own test), the + speculative config; the engine builds stubbed.""" + core = SimpleNamespace( + mapping=SimpleNamespace(tp_size=1, enable_attention_dp=False), + num_experts=128, + hidden=6144, + moe_layers=frozenset(range(1, num_layers)), + sparse_layers=frozenset(range(1, num_layers)), + attn=[_FakeSparseAttn(indexer_kv_dtype) for _ in range(num_layers)], + w={"embed": torch.empty(1, device="cuda")}, + mxfp8=SimpleNamespace(quant_method=SimpleNamespace(use_cutlass=True)), + model_config=SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=DRAFT_LEN, max_total_draft_tokens=DRAFT_LEN) + ), + moe_engines={}, + moe16_engines={}, + projections=lambda: [], + ) + for name in _CORE_METHODS: + setattr(core, name, getattr(modeling.MiniMaxM3Tp4ModelingV2, name).__get__(core)) + built = [] if built is None else built + + def build(i, num_tokens, shared=True): + if fail: + raise ValueError("the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI") + built.append(("moe", i, num_tokens, shared)) + return ("engine", i, num_tokens) + + def build16(i, num_tokens): + if fail: + raise ValueError("the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI") + built.append(("moe16", i, num_tokens)) + return ("engine16", i, num_tokens) + + core._build_engine, core._build_moe16_engine = build, build16 + return core + + +def _forbid_compiles(monkeypatch, core=None) -> None: + """Every compile / build entry point a forward could reach raises.""" + import cutlass.cute as cute + + def boom(*args, **kwargs): + raise AssertionError("a compile entry point ran after prebuild") + + monkeypatch.setattr(cute, "compile", boom) + monkeypatch.setattr(mx_gemv, "ensure_compiled", boom) + monkeypatch.setattr(mx_gemv, "precompile", boom) + monkeypatch.setattr(index_score_kernel.IndexScoreRunner, "_compile", boom) + monkeypatch.setattr(index_score_t2sc_kernel.IndexScoreRunner, "_compile", boom) + if core is not None: + core._build_engine = core._build_moe16_engine = boom + + +def test_decode_query_lens_follow_the_speculative_config() -> None: + core = _core() + assert core.decode_query_lens() == [1, 1 + DRAFT_LEN] + core.model_config.spec_config = None + assert core.decode_query_lens() == [1] + core.model_config.spec_config = SimpleNamespace( + max_draft_len=4, max_total_draft_tokens=4, draft_len_schedule={1: 4, 8: 2} + ) + assert core.decode_query_lens() == [1, 3, 5] + + +def test_prebuild_builds_every_engine_count_and_scorer_variant() -> None: + built = [] + core = _core(built=built) + times = core.prebuild_kernels() + + assert mx_gemv.frozen(), "prebuild must close the GEMV compiles" + + # Layer 0 is dense; every MoE layer gets every count, in layer order. + def per_layer(i): + engines = [("moe", i, n, True) for n in range(1, dispatch.MOE_ENGINE_MAX_TOKENS + 1)] + return engines + [ + ("moe16", i, n) for n in range(dispatch.MOE16_MIN_TOKENS, dispatch.MOE16_MAX_TOKENS + 1) + ] + + assert built == per_layer(1) + per_layer(2) + assert set(core.moe_engines[2]) == {(n, True, True) for n in dispatch.MOE_ENGINE.keys} + assert set(core.moe16_engines[2]) == set(dispatch.MOE16_ENGINE.keys) + + want = { + (torch.float8_e4m3fn, INDEX_HEADS, qlen, INDEX_DIM, modeling.INDEX_SCORE_SPLIT_K, early) + for qlen in (1, 1 + DRAFT_LEN) + for early in (True, False) + } + assert set(times["scorer_keys"]) == want + assert want <= set(index_score_kernel.IndexScoreRunner.kernel_cache) + for key in ("gemv_s", "moe_s", "scorer_s"): + assert times[key] >= 0.0 + + # The generation router's device decision, as the built-in gate makes it. + assert core.router_tc_cluster16 == bool(torch.ops.trtllm.minimax_m3_router_tc_cluster16(6144)) + + # A second build builds no engine again. + built.clear() + core.prebuild_kernels() + assert built == [] + + +def test_engine_lookups_after_prebuild_never_build(monkeypatch) -> None: + core = _core() + core._prebuild_engines(5) + _forbid_compiles(monkeypatch, core) + for require in (False, True): + dispatch.set_require(require) + for n in sorted(dispatch.MOE_ENGINE.keys): + assert core.engine(5, n, shared=True) == ("engine", 5, n) + for n in sorted(dispatch.MOE16_ENGINE.keys): + assert core.moe16_engine(5, n) == ("engine16", 5, n) + + +def test_missing_engine_is_a_miss_not_a_build(monkeypatch) -> None: + core = _core() + core.moe_engines[7], core.moe16_engines[7] = {}, {} + _forbid_compiles(monkeypatch, core) + dispatch.set_require(False) + assert core.engine(7, 3, shared=True) is None + assert core.moe16_engine(7, 10) is None + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + core.engine(7, 3, shared=True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + core.moe16_engine(7, 10) + + +def test_failed_engine_build_is_logged_then_missed() -> None: + core = _core(fail=True) + core._prebuild_engines(9) + assert all(v is None for v in core.moe_engines[9].values()) and len(core.moe_engines[9]) == 8 + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="build failed"): + core.engine(9, 1, shared=True) + + +def _projection(nk, push): + n, k = nk + w = (torch.randn(n, k, device="cuda") / k**0.5).to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return SimpleNamespace(weight=wq.view(n, k), weight_scale=wsf, push=push) + + +@pytest.mark.parametrize( + "nk,push", + ((O_PROJ, "attn"), (SHARED_DOWN, "moe"), (SHARED_GATE_UP, None)), + ids=("o_proj", "shared_down", "shared_gate_up"), +) +def test_gemvs_after_prebuild_never_compile(monkeypatch, nk, push) -> None: + use_pdl = True + module = _projection(nk, push) + variants = modeling._gemv_variants(*nk, push, 4) + pushes = [v for v in variants if v[2] is not None] + if push is None: + assert not pushes + else: + covered = dispatch.ATTN_PUSH.keys if push == "attn" else dispatch.MOE_PUSH.keys + assert {rows for rows, _, _ in pushes} == set(covered) + assert {p[1] for _, _, p in pushes} == {b for b in (8, 16) if b <= max(covered)} + mx_gemv.freeze(False) + # As prebuild_kernels compiles one projection. + mx_gemv.ensure_compiled(module.weight, module.weight_scale, use_pdl) + for rows, x_bf16, ar in variants: + mx_gemv.precompile(module.weight, module.weight_scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + mx_gemv.freeze(True) + _forbid_compiles(monkeypatch) + + n, k = nk + for rows, x_bf16, ar in variants: + if ar is not None: + # A push needs a multicast exchange to launch: its kernel is compiled (the forward finds it). + resolved = mx_gemv._resolve(rows, n, k, x_bf16, True, None, use_pdl, (0, 0, 0, *ar)) + if resolved is not None: + assert mx_gemv.is_compiled(resolved[0]), (rows, x_bf16, ar) + continue + if x_bf16: + x = torch.randn(rows, k, device="cuda").to(torch.bfloat16) + out = mx_gemv.m3_mx_gemv_bf16(x, module.weight, module.weight_scale, use_pdl=use_pdl) + else: + xq, xsf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(rows, k, device="cuda").to(torch.bfloat16), True + ) + out = mx_gemv.m3_mx_gemv( + xq.view(rows, k), xsf, module.weight, module.weight_scale, use_pdl=use_pdl + ) + served = mx_gemv.serves(n, k, rows, x_bf16) + assert (out is not None) == served, (rows, x_bf16) + assert mx_gemv.last_uncompiled is None + torch.cuda.synchronize() + + +def test_frozen_gemv_declines_an_uncompiled_kernel(monkeypatch) -> None: + module = _projection(SHARED_DOWN, None) + mx_gemv.freeze(True) + _forbid_compiles(monkeypatch) + # A launch option nothing prebuilt (no early trigger, without a push): declined, never compiled. + xq, xsf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(3, 768, device="cuda").to(torch.bfloat16), True + ) + key = mx_gemv._resolve(3, 6144, 768, False, False, None, True, None)[0] + if mx_gemv.is_compiled(key): + pytest.skip("this process already compiled the probe kernel") + out = mx_gemv.m3_mx_gemv( + xq.view(3, 768), xsf, module.weight, module.weight_scale, trigger_early=False, use_pdl=True + ) + assert out is None and mx_gemv.last_uncompiled == key + + +@pytest.mark.parametrize("indexer_kv_dtype", ("fp8", "bf16")) +def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> None: + core = _core(indexer_kv_dtype=indexer_kv_dtype) + keys = core._prebuild_index_scorer(core.decode_query_lens(), torch.device("cuda")) + assert len(keys) == 4 + _forbid_compiles(monkeypatch, core) + dtype = torch.float8_e4m3fn if indexer_kv_dtype == "fp8" else torch.bfloat16 + for _, heads, qlen, dim, split_k, early in keys: + rows = 1 if early else 2 + blocks = 3 + idx_q = torch.randn(rows * qlen, heads, dim, device="cuda").to(dtype) + idx_k = torch.randn(rows * blocks, 1, 128, dim, device="cuda").to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device="cuda").view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * 128 - 5, dtype=torch.int32, device="cuda") + score = torch.empty(heads, blocks, rows * qlen, dtype=torch.float32, device="cuda") + m3_index_score( + idx_q, + idx_k.squeeze(1), + block_table, + seq_lens, + score.transpose(1, 2), + qlen, + split_k, + early, + ) + torch.cuda.synchronize() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..cd7d800bd1ca --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode robustness without a GPU kernel: the multicast all-reduce epilogue options keep the cluster-barrier bit, +an engine build that fails on one TP rank drops that engine on every rank (two ranks simulated with threads), and +the Eagle3 fc fork's output is recorded on the stream that reads it.""" + +import threading +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = modeling.require_mode() + yield + modeling.set_require(require) + modeling.reset_logged() + + +@pytest.mark.parametrize("finalize", [False, True]) +@pytest.mark.parametrize("bucket", [8, 16, 32, 48]) +def test_ar_opts_keep_cluster_barrier_bit(monkeypatch, bucket, finalize): + # Tables without bit 2 and with buckets missing: every lookup, fallback included, keeps the barrier. + monkeypatch.setattr(mc, "AR_OPTS_BY_M", (0, 1, 4, 5, 8, 9, 12, 13)) + monkeypatch.setattr(mc, "AR_OPTS_BY_BUCKET", {}) + monkeypatch.setattr(mc, "AR_FIN_OPTS_BY_BUCKET", {}) + for m in range(1, min(bucket, 8) + 1): + opts = mc._ar_opts(m, bucket, finalize=finalize) + assert opts & mc._OPT_CLUSTER_BARRIER and 0 <= opts <= 15, (m, bucket, finalize, opts) + + +def test_ar_opts_shipped_tables_unchanged(): + for m in range(1, 9): + assert mc._ar_opts(m, 8) == mc.AR_OPTS_BY_M[m - 1] + for b, o in mc.AR_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b) == o + for b, o in mc.AR_FIN_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b, finalize=True) == o + + +class _FakeComm: + """An MPI-like TP communicator for ``world`` threads: allreduce sums the ranks' ints.""" + + def __init__(self, world): + self.world = world + self.barrier = threading.Barrier(world) + self.lock = threading.Lock() + self.vals = [] + self.calls = 0 + + def allreduce(self, v): + with self.lock: + self.vals.append(v) + self.calls += 1 + self.barrier.wait() + total = sum(self.vals) + self.barrier.wait() + with self.lock: + self.vals.clear() + self.barrier.wait() + return total + + def Get_size(self): # noqa: N802 (MPI name) + return self.world + + +def _mapping(tp_size=2, adp=False): + return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) + + +def _run_ranks(fns): + out, errs = [None] * len(fns), [] + + def run(i): + try: + out[i] = fns[i]() + except BaseException as e: # noqa: BLE001 + errs.append(e) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errs, errs + return out + + +@pytest.fixture +def two_rank_comm(monkeypatch): + from tensorrt_llm._torch.distributed import ops + + comm = _FakeComm(2) + monkeypatch.setattr(mc, "_tp_comm", lambda mapping: comm) + monkeypatch.setattr(ops, "mpi_disabled", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return comm + + +def test_tp_agree_single_rank_and_adp_take_no_collective(monkeypatch): + monkeypatch.setattr(mc, "_tp_comm", mock.Mock(side_effect=AssertionError("no collective"))) + assert mc.tp_agree(_mapping(1), True) is True + assert mc.tp_agree(_mapping(1), False) is False + assert mc.tp_agree(_mapping(4, adp=True), False) is False + + +def test_tp_agree_refuses_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError): + mc.tp_agree(_mapping(2), True) + + +@pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) +def test_tp_agree_built_two_ranks(two_rank_comm, oks): + built = [object() if ok else None for ok in oks] + out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + assert out == (built if all(oks) else [None, None]) + + +def _core(fail: bool) -> SimpleNamespace: + """The engine-building half of the core (``MiniMaxM3Tp4ModelingV2``) on one of two simulated TP ranks, its + engine builds stubbed (and failing on one rank).""" + core = SimpleNamespace(mapping=_mapping(2), moe_engines={}, moe16_engines={}) + cls = modeling.MiniMaxM3Tp4ModelingV2 + + def build(i, num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + core._build_engine = build + core._build_moe16_engine = lambda i, num_tokens: build(i, num_tokens) + for name in ("_prebuild_engines", "_try_build", "engine", "moe16_engine", "_lookup"): + setattr(core, name, getattr(cls, name).__get__(core)) + return core + + +@pytest.mark.parametrize("fail", [(False, False), (False, True)]) +def test_engine_build_agreed_across_ranks(two_rank_comm, fail): + cores = [_core(f) for f in fail] + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + keys = sorted(cores[0].moe_engines[3]) + sorted(cores[0].moe16_engines[3]) + assert keys and two_rank_comm.calls == 2 * len(keys) + for engines in ("moe_engines", "moe16_engines"): + a, b = getattr(cores[0], engines)[3], getattr(cores[1], engines)[3] + assert a.keys() == b.keys() + for key in a: + assert (a[key] is None) == (b[key] is None) == any(fail), (engines, key) + # Both ranks then take the same path in the forward: the engine, or (after a miss) the generic MoE. + modeling.set_require(False) + for n in sorted(modeling.MOE_ENGINE.keys): + got = [c.engine(3, n, shared=True) for c in cores] + assert (got[0] is None) == (got[1] is None) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + M3Eagle3Worker = modeling.M3Eagle3Worker + main, fork = mock.Mock(name="main"), mock.Mock(name="fork") + hidden = mock.Mock(name="hidden", is_cuda=True) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *a: main) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + worker = SimpleNamespace(_m3_fc_hidden=hidden) + M3Eagle3Worker._m3_fc_join(worker, fork) + main.wait_stream.assert_called_once_with(fork) + if capturing: + hidden.record_stream.assert_not_called() + else: + hidden.record_stream.assert_called_once_with(main) + worker._m3_fc_hidden = None + M3Eagle3Worker._m3_fc_join(worker, fork) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..1153eaee47bd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_silu_and_mul entry of the MiniMax-M3 tp8 target: SwiGLU-OAI and plain SwiGLU.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.activation.m3_silu_and_mul import ( # noqa: E501 + m3_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + # Gate and up spread past the clamp (|x| up to ~12) so both clamps engage. + x = (torch.randn(spec["m"], 2 * spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + return (x, *OAI) if spec["oai"] else (x, None, None, None) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few + fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); + measured on GB300: worst element 0.5 unit, relative RMS 0.01.""" + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(*args) + assert out.shape == (cell.spec["m"], cell.spec["inter"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=1.0, rms_ulp=0.1) + + +def test_the_gate_sees_a_missing_clamp() -> None: + """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" + spec = op.CELLS[0].spec + x, *oai = _build(dict(spec, m=64), 3) + out = op(x, *oai) + with pytest.raises(AssertionError): + assert_within_ulp(out, op.reference(x, None, oai[1], oai[2]), element_ulp=1.0, rms_ulp=0.1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_attn_runtime_tp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_attn_runtime_tp8.py new file mode 100644 index 000000000000..e9aaad5c65fa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_attn_runtime_tp8.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real MiniMax-M3 attention runtime state at the TP8 per-rank geometry, for the backend-step entries' GPU tests. + +One hybrid NVFP4 KV cache manager (``MiniMaxM3KVCacheManagerV2``) laid out as the target's: a dense layer in the FP8 +half, a sparse MSA layer in the NVFP4 half with its FP8 index-K cache, and (``draft=True``) the one-model Eagle3 +drafter's shared layer, read through the manager's FP8 draft view as the runtime's draft swap reads it. Metadata is +the runtime's own type, filled by its own ``prepare()``; the backends are built the way the target builds them. + +The geometry is one rank's: 8 query heads over one KV head (M3's 4 KV heads, each duplicated on a rank pair) and +one index head of dim 128. The mapping is a single rank holding exactly those heads, so nothing is sharded again. + +History rows are written by the in-tree scatters (``fused_write_layer_caches`` for the FP8 layers, +``fused_write_layer_caches_nvfp4`` for the NVFP4 layer and its index-K), at the slots the manager assigned, which are +not in request order. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + MiniMaxM3MsaSparseAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, + fused_write_layer_caches_nvfp4, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HEADS = 8 +KV_HEADS = 1 +INDEX_HEADS = 1 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +DENSE, SPARSE, DRAFT = 0, 1, 2 +MAX_BATCH = 8 +MAX_SEQ_LEN = 8192 +MAX_NUM_TOKENS = 8192 +#: The Eagle3 drafter checkpoint's RoPE: NeoX over the whole head, theta 5e6. +DRAFT_ROPE_THETA = 5.0e6 + +#: E4M3-exact planted index levels, 1/8 apart in relative terms (0.25 .. 15): 48 distinct block scores. +LEVELS = tuple(2.0**e * (1 + m / 8) for e in range(-2, 4) for m in range(8)) +#: The index value of a block's non-peak rows; its score (8) is below every level's (>= 16). +LOW = 0.125 + + +def mapping() -> Mapping: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + +def sparse_attention_config() -> MiniMaxM3SparseAttentionConfig: + """The target's runtime: MSA, the FP8 indexer, the fused five-way projection (M3's checkpoint selection rule: + no forced initial blocks, the query's own block forced).""" + return MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_num_index_heads=INDEX_HEADS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + ) + + +def make_manager(*, draft: bool) -> MiniMaxM3KVCacheManagerV2: + torch.cuda.init() + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_tokens=4 * MAX_SEQ_LEN, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + event_buffer_max_size=0, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=[KV_HEADS] * (3 if draft else 2), + head_dim=HEAD_DIM, + tokens_per_block=PAGE, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + mapping=mapping(), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") if draft else None + ), + sparse_layer_ids=[SPARSE], + sparse_attention_config=sparse_attention_config(), + ) + + +def make_msa_backend(layer: int) -> MiniMaxM3MsaSparseAttention: + """One layer's MSA backend as ``MiniMaxM3Tp8ModelingV2`` builds it (NVFP4 KV quant config on every layer).""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + backend = create_attention( + "TRTLLM", + layer, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=None, + quant_config=quant, + q_scaling=1.0, + sparse_params=sparse_attention_config().to_sparse_params(layer_idx=layer), + ) + backend.update_quant_config(quant) + assert isinstance(backend, MiniMaxM3MsaSparseAttention), type(backend).__name__ + return backend + + +def make_draft_backend() -> TrtllmAttention: + """The drafter layer's attention as ``M3Eagle3Draft`` builds it: fused NeoX RoPE, an FP8 KV cache (the target + overrides the drafter's KV dtype to FP8), no other quantization.""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + backend = create_attention( + "TRTLLM", + DRAFT, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=HEAD_DIM, theta=DRAFT_ROPE_THETA, max_positions=MAX_SEQ_LEN), + ), + quant_config=quant, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + backend.update_quant_config(quant) + assert type(backend) is TrtllmAttention, type(backend).__name__ + return backend + + +def make_msa_metadata(manager) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = MiniMaxM3MsaSparseAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + max_num_sequences=MAX_BATCH, + kv_cache_manager=manager, + mapping=mapping(), + sparse_metadata_params=sparse_attention_config().to_sparse_metadata_params(), + ) + # The target's MSA preparation choice that changes what prepare() stages (the identity tables of short steps). + metadata.msa_policy = MsaRuntimePolicy(short_sequence_identity=True) + return metadata + + +def make_draft_metadata(manager) -> TrtllmAttentionMetadata: + """Drafter metadata over the shared manager's FP8 draft view, the manager the runtime's draft swap installs.""" + view = manager.get_draft_kv_cache_view() + assert view is not None and view.dtype == DataType.FP8 + return TrtllmAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=mapping(), + ) + + +def allocate(manager, request_ids: Sequence[int], qlens: Sequence[int], cached: Sequence[int]): + manager.add_dummy_requests(list(request_ids), token_nums=[c + n for c, n in zip(cached, qlens)]) + + +def prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + c + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +def flat_slots(manager, layer: int, request_id: int, positions: torch.Tensor) -> torch.Tensor: + """The cache slots (page slot id * PAGE + offset) of a request's positions in ``layer``.""" + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0], + dtype=torch.int64, + device="cuda", + ) + positions = positions.to(device="cuda", dtype=torch.int64) + return (pages[positions // PAGE] * PAGE + positions % PAGE).to(torch.int32) + + +def write_fp8_rows(manager, layer, request_id, positions, k, v) -> None: + """K / V rows [n, 128] (E4M3) into an FP8 layer's cache at ``positions``.""" + view = manager.get_buffers(layer, "HND") + assert view.dtype == torch.float8_e4m3fn + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches(view[:, 0], view[:, 1], None, slots, k, v, None) + + +def write_nvfp4_rows(manager, metadata, layer, request_id, positions, k, v, idx_k, quant) -> None: + """BF16 K / V rows quantized into the NVFP4 layer's cache, and index-K rows into its FP8 index cache.""" + data = manager.get_buffers(layer, "HND") + scales = manager.get_block_scale_buffers(layer, "HND") + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches_nvfp4( + data[:, 0], + data[:, 1], + scales[:, 0], + scales[:, 1], + metadata.msa_idx_k_cache(layer), + slots, + k, + v, + idx_k, + quant, + ) + + +def index_directions(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """Two orthogonal ±1 index directions (the halves of the 128 dims), E4M3-exact.""" + signs = torch.randint(0, 2, (HEAD_DIM,), generator=gen, device="cuda").float() * 2 - 1 + u0 = torch.zeros(HEAD_DIM, device="cuda") + u1 = torch.zeros(HEAD_DIM, device="cuda") + u0[: HEAD_DIM // 2] = signs[: HEAD_DIM // 2] + u1[HEAD_DIM // 2 :] = signs[HEAD_DIM // 2 :] + return u0, u1 + + +def planted_index_keys(length: int, u0, u1, gen: torch.Generator) -> torch.Tensor: + """Index-K rows [length, 128] (bf16, E4M3-exact) whose per-block maximum score is unambiguous. + + Every block has one peak row at a random position, ``a_b u0 + c_b u1`` with ``a`` and ``c`` two independent + permutations of ``LEVELS``; its other rows are ``LOW (u0 + u1)``. A query along ``u0`` then ranks the blocks by + ``a``, one along ``u1`` by ``c``, every score 1/8 or more apart from the next.""" + blocks = (length + PAGE - 1) // PAGE + assert blocks <= len(LEVELS) + levels = torch.tensor(LEVELS, device="cuda") + a = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + c = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + rows = (LOW * (u0 + u1)).expand(length, HEAD_DIM).clone() + for b in range(blocks): + span = min(PAGE, length - b * PAGE) + peak = b * PAGE + int(torch.randint(0, span, (1,), generator=gen, device="cuda")) + rows[peak] = a[b] * u0 + c[b] * u1 + return rows.to(torch.bfloat16) + + +def nvfp4_quant(*values: torch.Tensor) -> torch.Tensor: + """The (Q, K, V) quantization scales ``fused_write_layer_caches_nvfp4`` takes: each role's amax at the top of the + E2M1 x E4M3 range.""" + quant = torch.ones(3, dtype=torch.float32, device="cuda") + for role, x in enumerate(values, start=1): + quant[role] = (448.0 * 6.0) / x.abs().max().float() + return quant diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_nvfp4_cache.py new file mode 100644 index 000000000000..bfa03c09da03 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_nvfp4_cache.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A MiniMax-M3 TP8 sparse-layer decode step over a packed NVFP4 cache, for the attention entries' tests. + +The cache is written by the in-tree scatter (``fused_write_layer_caches_nvfp4``, +which quantizes with ``torch.ops.trtllm.fp4_quantize``), into one +``[pages, 2 (K, V), 1, 128, .]`` pool per role exactly as ``msa_paged_kv`` and +``get_block_scale_buffers`` hand the K and V views over. Pages are shuffled so +a kernel that ignored the block table would read the wrong ones. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +#: Query heads per rank at TP8 (on one KV head). +NUM_HEADS = 8 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 + + +def token_kv_len(seq_len: int, decode_query_len: int, offset: int) -> int: + """The KV extent query token ``offset`` of a request attends, as every decode kernel derives it.""" + return max(seq_len - decode_query_len + offset + 1, 0) + + +def build_step( + seq_lens: Sequence[int], decode_query_len: int, *, identity: bool, seed: int +) -> SimpleNamespace: + """One decode step: E4M3 q [T, 8, 128], the packed NVFP4 K/V views and their scales, a [1, T, 16] + head-major top-k table (the identity table, or a random ascending selection with a -1 tail), the block + table and per-request lengths.""" + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = max(1, max((s + PAGE - 1) // PAGE for s in seq_lens)) + num_pages = batch * max_blocks + num_slots = num_pages * PAGE + block_table = ( + torch.randperm(num_pages, device="cuda", generator=g) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + data_pool = torch.zeros(num_pages, 2, 1, PAGE, HEAD_DIM // 2, dtype=torch.uint8, device="cuda") + scale_pool = torch.zeros( + num_pages, 2, 1, PAGE, HEAD_DIM // 16, dtype=torch.uint8, device="cuda" + ) + quant = torch.ones(3, dtype=torch.float32, device="cuda") + values = [] + for role in (0, 1): + x = torch.randn(num_slots, HEAD_DIM, device="cuda", generator=g, dtype=torch.float32).to( + torch.bfloat16 + ) + quant[role + 1] = (448.0 * 6.0) / x.abs().max().float().item() + values.append(x) + assert fused_write_layer_caches_nvfp4( + data_pool[:, 0], + data_pool[:, 1], + scale_pool[:, 0], + scale_pool[:, 1], + None, + torch.arange(num_slots, device="cuda", dtype=torch.int32), + values[0], + values[1], + None, + quant, + ) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) + topk = torch.full((1, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + nblocks = (token_kv_len(seq_lens[req], decode_query_len, offset) + PAGE - 1) // PAGE + if identity: + assert nblocks <= TOPK, "an identity step has every request within 16 blocks" + chosen = torch.arange(nblocks, device="cuda") + else: + chosen = torch.randperm(nblocks, device="cuda", generator=g)[:TOPK].sort().values + topk[0, tok, : chosen.numel()] = chosen.to(torch.int32) + return SimpleNamespace( + q=q, + k_paged=data_pool[:, 0], + v_paged=data_pool[:, 1], + k_block_scale=scale_pool[:, 0], + v_block_scale=scale_pool[:, 1], + k_global_scale=(1.0 / quant[1:2]).contiguous(), + v_global_scale=(1.0 / quant[2:3]).contiguous(), + topk=topk, + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device="cuda"), + decode_query_len=decode_query_len, + sm_scale=HEAD_DIM**-0.5, + source=values, + quant=quant, + num_slots=num_slots, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..bef0db5ce41b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7): the TP8 dense layers' FP8 q/k/v (minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8, head-count +generic, at the TP8 per-rank 8 query heads over 1 KV head) carries the bits of the stock path, the BF16 +fused_qk_norm_rope output cast to E4M3, for every norm / RoPE variant: Q and K are rounded to BF16 before E4M3 (K is +stored in the FP8 cache), V is cast.""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers the trtllm ops) + +assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) + +HQ, HK, HV, D, ROT = 8, 1, 1, 128, 64 + + +@pytest.mark.parametrize("tokens", [1, 4, 8, 4096]) +@pytest.mark.parametrize("use_gemma", [False, True], ids=["rmsnorm", "gemma"]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_dense_fp8_qkv_is_the_bf16_kernel_cast(tokens, use_gemma, is_neox) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8") + g = torch.Generator(device="cuda").manual_seed(tokens * 4 + int(use_gemma) * 2 + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + ours = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8( + qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, use_gemma + ) + stock = qkv.clone() + torch.ops.trtllm.fused_qk_norm_rope( + stock, HQ, HK, HV, D, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, + 1.0, 0.0, 0.0, 1.0, True, use_gemma, False, 0, 0, + ) # fmt: skip + assert torch.equal(ours.view(torch.uint8), stock.to(torch.float8_e4m3fn).view(torch.uint8)) + + +@pytest.mark.parametrize("tokens,live", [(1, 1), (4, 4), (8, 5), (64, 64)]) +@pytest.mark.parametrize("tokens_per_block", [32, 128]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_fused_kv_insert_stores_the_bytes_of_write_kv_slots_and_the_compact_q( + tokens, live, tokens_per_block, is_neox +) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + g = torch.Generator(device="cuda").manual_seed(tokens * 1000 + tokens_per_block + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + pages, layers, layer = 16, 3, 1 + # The layer's HND view of a multi-layer pool, strided as get_buffers returns it. + raw = torch.randint( + 0, + 256, + (pages, layers, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.randperm(pages * tokens_per_block, device="cuda", generator=g)[:tokens].to( + torch.int32 + ) + slots[live:] = -1 # padded rows own no slot + + def view(t): + return t.view(torch.float8_e4m3fn)[:, layer].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, False) + q_compact = torch.empty(tokens, HQ, D, device="cuda", dtype=torch.float8_e4m3fn) + ours = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, live, q_compact + ) + plain = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + assert torch.equal(ours.view(torch.uint8), plain.view(torch.uint8)) + assert torch.equal( + q_compact.view(tokens, HQ * D).view(torch.uint8), plain[:, : HQ * D].view(torch.uint8) + ) + + ref_raw = raw.clone() + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D) + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D) + ref = view(ref_raw) + write_kv_slots(ref[:, 0], slots, k, live, layout="HND") + write_kv_slots(ref[:, 1], slots, v, live, layout="HND") + assert torch.equal(ours_raw, ref_raw) + + +def test_the_fused_kv_insert_skips_a_negative_slot() -> None: + """A live row whose slot is the -1 sentinel stores nothing (no write before the pool row).""" + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + tokens, tokens_per_block, pages = 4, 32, 8 + g = torch.Generator(device="cuda").manual_seed(7) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + raw = torch.randint( + 0, + 256, + (pages, 3, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.tensor([5, -1, 40, 77], device="cuda", dtype=torch.int32) + + def view(t): + return t.view(torch.float8_e4m3fn)[:, 1].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, True, pos, False) + torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, tokens + ) + plain = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + keep = torch.tensor([0, 2, 3], device="cuda") + ref_raw = raw.clone() + ref = view(ref_raw) + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D)[keep] + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D)[keep] + write_kv_slots(ref[:, 0], slots[keep], k, 3, layout="HND") + write_kv_slots(ref[:, 1], slots[keep], v, 3, layout="HND") + assert torch.equal(ours_raw, ref_raw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..c89183995648 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_attention entry: the Eagle3 drafter layer's attention against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp8``): the target's hybrid NVFP4 +``MiniMaxM3KVCacheManagerV2`` with its one-model Eagle3 draft layer, the requests allocated in it, drafter +``TrtllmAttentionMetadata`` over the manager's FP8 draft view (the manager the runtime's draft swap installs) filled +by its ``prepare()``, and the drafter's ``TrtllmAttention`` built as ``M3Eagle3Draft`` builds it (fused NeoX RoPE at +the drafter checkpoint's theta, FP8 KV). Earlier rows go into the draft layer's cache through the in-tree scatter at +the slots the manager assigned; the call rotates and appends the step's own. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp8 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. +# The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: +# it takes the rotated Q in E4M3 (rounding the reference's Q to E4M3 takes the RMS gap from ~9 to ~7 ulps) and rounds +# the softmax probabilities before the PV product. Measured on GB300 over these cells (max element / RMS): 14.8 / 8.7 +# (4 x 4 rows), 13.3 / 7.4 (8 x 1), 19.1 / 7.7 (512 context rows). The band leaves about 1.5x; a reference shifted +# by one row is far outside it. +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 + +# Cached tokens per request of each cell; the new rows per request come from the cell. +_CACHED = { + (2, 4): (200, 1000), + (2, 1): (127, 2047), + (1, None): (0,), +} +#: Rows past the live ones in the fused QKV, as a padded step hands them over. +_PAD = 3 + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + +# Non-paged context FMHA on SM107 does not apply RoPE (this fixture's path; the served drafter runs the paged context +# FMHA, which rotates): the context cell is a strict expected failure until the runtime rotates there too. +_UNROTATED_CONTEXT = pytest.mark.xfail( + strict=True, + reason="non-paged context FMHA on SM107 does not apply RoPE (shared runtime)", +) +_CELLS = [ + pytest.param(c, marks=_UNROTATED_CONTEXT) if c.spec.get("context") else c for c in op.CELLS +] + + +@pytest.mark.parametrize("cell", _CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + context = spec.get("context") + cached = _CACHED[spec["reqs"], spec.get("qlen")] + qlens = (context,) if context else (spec["qlen"],) * spec["reqs"] + assert len(cached) == spec["reqs"] + contexts = 1 if context else 0 + rids = list(range(spec["reqs"])) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = rt.make_draft_metadata(manager) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + assert int(metadata.num_tokens) == total + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + certifying(op) + out = op(backend, qkv, metadata) + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{cell.why}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + group = rt.HEADS // rt.KV_HEADS + for h in range(rt.KV_HEADS): + cols = slice(h * group * rt.HEAD_DIM, (h + 1) * group * rt.HEAD_DIM) + element, rms = _ulp_stats(out[:, cols], ref[:, cols]) + print(f"[{cell.why}] KV head {h}: max element {element:.2f} ulp, rms {rms:.3f} ulp") + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + # The call's append: each live row's rotated K and its V at the row's position, E4M3 (unit KV scale), to + # within one E4M3 step of the float64 rotation (the kernel rotates in FP32). + pool = manager.get_buffers(rt.DRAFT, "HND") + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + for h in range(rt.KV_HEADS): + k_cols = q_size + h * rt.HEAD_DIM + v_cols = q_size + (rt.KV_HEADS + h) * rt.HEAD_DIM + want_k = _rope(rows[:, k_cols : k_cols + rt.HEAD_DIM].float(), positions) + want_v = rows[:, v_cols : v_cols + rt.HEAD_DIM].to(torch.float8_e4m3fn) + got_k = pool[slots // rt.PAGE, 0, h, slots % rt.PAGE].to(torch.float64) + got_v = pool[slots // rt.PAGE, 1, h, slots % rt.PAGE] + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), ( + f"K head {h} is appended rotated, in E4M3" + ) + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)), f"V head {h}" + first += n + + # The band discriminates: a reference whose rows are shifted by one token is outside it. + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py new file mode 100644 index 000000000000..6fa19e3352fd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for kernels/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V +append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row +first draft pass and the 1-row later passes of 1-2 requests. + +Lengths sweep the KV tile edges each query row in turn ends on (the rows of a Q CTA with no key in its trailing +splits), with mixed per-request lengths; every call is held to the fp32 reference over the cache after the append +(E4M3 query, unit KV scales, no RoPE) within the FP8-KV band, finite, for both split-KV reducers. +""" + +import math +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) + +assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +FP8_TOL = 6e-2 +# Persistent like the FMHA manager's workspace and the spec-decoding metadata of a real step: the in-tree preprocess +# kernels must not see these move between calls. +_PERSISTENT = {} + + +def _persistent(name: str, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _edge_lens(plan, qlen: int): + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +def _library(qlen: int, nreq: int, own_reducer: bool): + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = own_reducer + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = torch.cuda.get_device_properties(0).multi_processor_count + lib._on_taken = lambda rows: None + # The in-tree preprocess pass (its RoPE-free cache append); the one-launch preprocess has its own bitwise test. + lib._fused_rope_mode = None + lib._pdl = False + lib._on_rope_taken = None + return lib + + +def _attn(qlen: int): + return SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=0, theta=10000.0, scale_type=0, scale=1.0, max_positions=1024 + ), + ) + + +def _run(lib, qlen: int, kv_lens, seed: int, spec: bool = True): + """One drafter generation pass over a fresh FP8 cache; returns (output, fp32 reference). ``spec`` passes the + spec-decoding lengths/positions (a multi-token generation step); without them a multi-row pass is the drafter's + plain generation step of input_seq_length rows, as it runs in-model.""" + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g).clamp_(-4, 4) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + for r, n in enumerate( + kv_lens + ): # shift the tail third of V, so a missed trailing split fails the band + start = 2 * n // 3 + for pg in range(start // PAGE, math.ceil(n / PAGE)): + pool[int(perm[r, pg]), 1, :, (start % PAGE if pg == start // PAGE else 0) :] += 1.0 + pool = pool.to(torch.float8_e4m3fn) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + pointers = torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64) + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=dev, generator=g) + seq = torch.tensor(kv_lens, dtype=torch.int32, device=dev) + one = torch.ones(1, dtype=torch.float32, device=dev) + multi = qlen > 1 and spec + out = torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=_attn(qlen), + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=pointers, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=seq, + spec_decoding_generation_lengths=( + _persistent( + f"gen_lengths{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=( + _persistent( + f"position_offsets{qlen}x{nreq}", + lambda: ( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous() + ), + ) + if multi + else None + ), + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=out, + ) + lib.run_generation(params) + torch.cuda.synchronize() + # fp32 reference over the cache after the append: E4M3 query, unit scales, causal over the step's rows. + q = qkv[:, : H * D].float().view(nreq, qlen, H, D).to(torch.float8_e4m3fn).float() + flat = pool.float() + ref = torch.empty(nreq, qlen, H, D, device=dev) + for r, n in enumerate(kv_lens): + blocks = perm[r, : math.ceil(n / PAGE)] + k = flat[blocks, 0].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + v = flat[blocks, 1].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + k = k.repeat_interleave(H // KVH, dim=0) + v = v.repeat_interleave(H // KVH, dim=0) + s = torch.einsum("qhd,hkd->hqk", q[r], k) * D**-0.5 + visible = ( + torch.arange(n, device=dev)[None, :] + < (n - qlen + 1 + torch.arange(qlen, device=dev))[:, None] + ) + s = s.masked_fill(~visible[None], float("-inf")) + ref[r] = torch.einsum("hqk,hkd->qhd", torch.softmax(s, dim=-1), v) + return out.float().view(nreq, qlen, H, D), ref + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_the_drafter_pass_at_tile_edges(qlen, nreq) -> None: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if fullctx.has_own_reducer() else (False,) + for own in reducers: + fullctx.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_LEN, own, get_env_enable_pdl() + ) + lens = _edge_lens(fullctx.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + lens = [n for n in lens if n > qlen] + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for i, kv_lens in enumerate(batches): + # The first draft pass runs in-model as a plain multi-row generation step; every other batch also takes the + # spec-decoding form. + spec = qlen > 1 and i % 2 == 1 + for own in reducers: + got, ref = _run(_library(qlen, nreq, own), qlen, kv_lens, sum(kv_lens), spec=spec) + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +def _drafter_model(scale: float): + """A drafter whose one attention module holds a TrtllmAttention backend as ``attn`` (the backend is no module).""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + backend = TrtllmAttention.__new__(TrtllmAttention) + backend.kv_cache_scaling_factor = torch.full((1,), scale, device="cuda") + backend.is_mla_enable, backend.kv_lora_rank, backend.v_head_dim, backend.head_dim = ( + False, + None, + None, + D, + ) + backend._fmha_manager = SimpleNamespace(fmha_libs=[]) + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.zeros(1, device="cuda")) + model.self_attn = torch.nn.Module() + model.self_attn.attn = backend + return model, backend + + +def test_install_declines_a_scaled_kv_cache() -> None: + """At unit KV scales the arithmetic is the stock decode's; a scaled cache keeps the stock path.""" + reasons = [] + model, layer = _drafter_model(0.5) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 0 + assert layer._fmha_manager.fmha_libs == [] and reasons == ["a scaled FP8 KV cache"] + model, layer = _drafter_model(1.0) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 1 + assert isinstance(layer._fmha_manager.fmha_libs[0], drafter.M3DrafterFullCtxFmha) + + +def test_a_cache_bound_past_the_plan_capacity_keeps_the_stock_path() -> None: + """The plans are compiled for ``max_kv_len`` of K/V: a step whose cache length bound is longer falls through.""" + + from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase + + class _Backend: + num_heads, num_kv_heads, head_dim, quant_mode = H, KVH, D, FP8_KV_QUANT_MODE + kv_scale_orig_quant = kv_scale_quant_orig = None + + backend = _Backend() + backend.kv_scale_orig_quant, backend.kv_scale_quant_orig = object(), object() + reasons = [] + lib = _library(1, 1, False) + lib._on_uncovered = reasons.append + lib._attn_ref = weakref.ref(backend) + fwd = SimpleNamespace( + attention_input_type=None, + kv_scale_orig_quant=backend.kv_scale_orig_quant, + kv_scale_quant_orig=backend.kv_scale_quant_orig, + ) + meta = SimpleNamespace(num_contexts=0, num_generations=1, max_seq_len=MAX_LEN + 8) + q = torch.empty(1, H * D, device="cuda") + assert not lib._is_supported(q, None, None, meta, fwd, phase=FmhaPhase.GENERATION) + assert reasons == ["the cache length bound is beyond the plans' K/V capacity"] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py new file mode 100644 index 000000000000..03fc2b3576b3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -0,0 +1,371 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for minimax_m3_sm107tp8_drafter_rope_fp8_kv, the TP8 drafter layer's one-launch generation +preprocess: bitwise against the in-tree ``thop.trtllm_gen_generation_preprocess`` followed by the torch E4M3 cast +of its BF16 query, on the query bytes and on every byte of the paged cache, with GPT-NeoX RoPE from a cos/sin table +(full and partial rotary width), multi-token (spec-decoding positions) and plain passes, one and two requests, and +query values past the E4M3 range. The library's fused pass is also held bitwise to its preprocess pass end to end. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp8_drafter_rope_fp8_kv"): + pytest.skip("minimax_m3_sm107tp8_drafter_rope_fp8_kv is not built", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +ROPE_GPT_NEOX = 2 +_PERSISTENT = {} + + +def _persistent(name, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _case(qlen, kv_lens, rotary_dim, multi, seed): + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = (torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g) * 2).to( + torch.float8_e4m3fn + ) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, device=dev, generator=g) * 4 + # A few query / key values near and past the E4M3 range (both casts saturate at 448). + hot = torch.rand(qkv.shape, device=dev, generator=g) < 0.01 + qkv = torch.where( + hot, qkv.sign() * (430 + 80 * torch.rand(qkv.shape, device=dev, generator=g)), qkv + ) + qkv = qkv.to(torch.bfloat16).contiguous() + max_pos = max(kv_lens) + 64 + cos_sin = torch.randn(max_pos, rotary_dim // 2, 2, device=dev, generator=g).float().contiguous() + attn = SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=ROPE_GPT_NEOX, + rotary_inv_freq=None, + rotary_cos_sin=cos_sin, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=rotary_dim, theta=10000.0, scale_type=0, scale=1.0, max_positions=max_pos + ), + ) + one = torch.ones(1, dtype=torch.float32, device=dev) + spec_offsets = None + if multi: + spec_offsets = _persistent( + f"pos{qlen}x{nreq}", + lambda: torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous(), + ) + # Spec positions are not the plain row index: shift them per request. + spec_offsets.copy_( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1) + + torch.arange(nreq, dtype=torch.int32, device=dev)[:, None] + ) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=attn, + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=torch.tensor(kv_lens, dtype=torch.int32, device=dev), + spec_decoding_generation_lengths=( + _persistent( + f"len{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=spec_offsets, + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev), + ) + return params, pool + + +def _preprocess(params, mpc): + """run_generation's in-tree preprocess call and the E4M3 cast of its query.""" + attn, meta, fwd, rope = params.attn, params.meta, params.fwd, params.attn.rope_params + q_processed = thop.trtllm_gen_generation_preprocess( + params.qkv_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope.dim, + rope.theta, + int(rope.scale_type), + rope.scale, + rope.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + False, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + mpc, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + )[0] + rows = params.num_tokens * attn.num_heads * attn.head_dim + return ( + q_processed.flatten()[:rows] + .to(torch.float8_e4m3fn) + .view(params.num_tokens, attn.num_heads, attn.head_dim) + ) + + +def _op(params, pool, mode): + attn = params.attn + nreq, qlen = params.batch_size, params.input_seq_length + kv_pages = pool.view(-1, KVH, PAGE, D) + tables = params.meta.kv_cache_block_offsets[0, :nreq] + q = torch.empty( + params.num_tokens, + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=pool.device, + ) + multi = params.spec_decoding_generation_lengths is not None and qlen > 1 + torch.ops.trtllm.minimax_m3_sm107tp8_drafter_rope_fp8_kv( + params.qkv_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + attn.rope_params.dim, + PAGE, + mode, + True, + ) + return q + + +# One-request multi-token passes keep the in-tree preprocess (it reads an unbuilt token map), so the spec cases hold +# two requests. +CASES = [ + (qlen, kv_lens, rotary, multi) + for qlen, multi in ((1, False), (4, True), (4, False)) + for kv_lens in (((4, 130) if multi else (4,)) if qlen > 1 else (130,), (8191, 131073)) + for rotary in (128, 64) +] + + +def _bytes(t): + return t.contiguous().view(torch.uint8) + + +def _matches(case, seed): + """Per rope mode: (query bytes equal, cache bytes equal) against the in-tree preprocess + cast.""" + qlen, kv_lens, rotary, multi = case + kv_lens = [max(n, qlen) for n in kv_lens] + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(qlen, kv_lens, rotary, multi, seed) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + pool_ref = pool.clone() + out = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + out[mode] = ( + bool(torch.equal(_bytes(q), _bytes(q_ref))), + bool(torch.equal(_bytes(pool), _bytes(pool_ref))), + ) + return out + + +@pytest.mark.parametrize( + "case", + CASES, + ids=lambda c: f"q{c[0]}-{'x'.join(map(str, c[1]))}-r{c[2]}-{'spec' if c[3] else 'plain'}", +) +def test_the_pinned_rope_mode_is_bitwise(case) -> None: + result = _matches(case, seed=7) + print(f"[rope modes] {case}: {result}", flush=True) + assert result[modeling.DRAFTER_ROPE_MODE] == (True, True), result + + +def test_the_fused_pass_matches_the_preprocess_pass() -> None: + mpc = torch.cuda.get_device_properties(0).multi_processor_count + for qlen, nreq in sorted(modeling.FULLCTX_DECODE_RECIPES): + for multi in (False, True) if qlen > 1 else (False,): + kv_lens = [70001 + 517 * r for r in range(nreq)] + outs = [] + for fused in (False, True): + params, pool = _case(qlen, kv_lens, 128, multi, seed=11) + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = True + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = mpc + lib._fused_rope_mode = modeling.DRAFTER_ROPE_MODE if fused else None + lib._pdl = True + taken = [] + lib._on_rope_taken = lambda rows, what: taken.append(what) + lib._on_taken = lambda rows: None + fullctx.prebuild(torch.device("cuda"), lib._recipes, MAX_LEN, True, True) + lib.run_generation(params) + torch.cuda.synchronize() + assert bool(taken) == (fused and not (multi and nreq == 1)) + outs.append((params.output.clone(), pool.clone())) + assert torch.equal(_bytes(outs[0][1]), _bytes(outs[1][1])), (qlen, nreq, multi, "cache") + assert torch.equal(_bytes(outs[0][0]), _bytes(outs[1][0])), ( + qlen, + nreq, + multi, + "output", + ) + + +def _round_bf16(x): + u = x.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000 + return u.to(torch.int64).to(torch.int32).view(torch.float32) + + +def _discriminating(n, seed): + """(cos, sin, x, pair) quadruples whose rotated first-half output's E4M3 byte differs between the rope modes + (fp32 CPU emulation; the GPU comparison is the verdict).""" + g = torch.Generator().manual_seed(seed) + keep = [] + while sum(k.shape[0] for k in keep) < n: + m = 1 << 20 + c, s = torch.randn(m, generator=g), torch.randn(m, generator=g) + x = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + xp = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + sn = -s + r0 = (c * x) + (sn * xp) + r1 = ((c.double() * x.double()) + (sn * xp).double()).float() + r2 = ((sn.double() * xp.double()) + (c * x).double()).float() + b = [ + _round_bf16(r.contiguous()).to(torch.float8_e4m3fn).view(torch.uint8) + for r in (r0, r1, r2) + ] + diff = (b[0] != b[1]) | (b[0] != b[2]) | (b[1] != b[2]) + diff &= (r0.abs() < 400) & (r1.abs() < 400) + keep.append(torch.stack([c, s, x, xp], 1)[diff]) + return torch.cat(keep)[:n] + + +def test_crafted_inputs_pin_the_rope_mode() -> None: + """Rotations whose E4M3 byte depends on how the two products are summed: the in-tree preprocess and exactly the + pinned mode agree on every one.""" + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(1, [130], 128, False, seed=3) + half = 64 + quads = _discriminating(H * half, seed=5).cuda() + pos = 129 # the one row's rope position: cache length 130, one query row + table = params.attn.rotary_cos_sin.view(-1, half, 2) + # One (cos, sin) per dim: every head shares the position's table row, so the heads take the same dims' quads with + # their own (x, pair). + table[pos, :, 0] = quads[:half, 0] + table[pos, :, 1] = quads[:half, 1] + qkv = params.qkv_input.view(-1, H + 2 * KVH, D) + for h in range(H): + qkv[0, h, :half] = quads[:half, 2].to(torch.bfloat16) + qkv[0, h, half:] = quads[:half, 3].to(torch.bfloat16) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + agree = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + agree[mode] = int((_bytes(q)[0, :, :half] == _bytes(q_ref)[0, :, :half]).sum()) + print( + f"[rope modes crafted] agreeing first-half query bytes of {H * half}: {agree}", flush=True + ) + assert agree[modeling.DRAFTER_ROPE_MODE] == H * half, agree diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..2179f1720e0d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp8_indexer_qk_norm_rope entry of the MiniMax-M3 tp8 target: index-Q out and index-K in the +cache against the in-place fused norm + RoPE followed by the E4M3 cast, byte for byte up to the 0.1% budget the +in-tree test of the kernel allows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 + m3_fp8_indexer_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" + +DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 + + +def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: + assert a.shape == b.shape and a.dtype == b.dtype == torch.float8_e4m3fn + differ = int((a.view(torch.uint8) != b.view(torch.uint8)).sum()) + assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + m, heads = cell.spec["m"], cell.spec["heads"] + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + qk = torch.randn(m, (heads + 1) * DIM, generator=g, device="cuda").to(torch.bfloat16) + qw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (m,), generator=g, device="cuda", dtype=torch.int32) + pages = (m + PAGE - 1) // PAGE + 1 + cache = torch.zeros(pages, 1, PAGE, DIM, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(pages * PAGE, generator=g, device="cuda")[:m].to(torch.int32) + before = qk.clone() + args = (qk, cache, slots, heads, DIM, ROTARY, EPS, qw, kw, THETA, pos) + with validating(op): + q = op(*args) + assert torch.equal(qk, before), "the BF16 input is left untouched" + q_ref, k_ref = op.reference(*args) + _bytes_close(q.reshape(m, heads, DIM), q_ref) + flat = cache.view(-1, DIM) + _bytes_close(flat[slots.long()], k_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_saturate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_saturate.py new file mode 100644 index 000000000000..900e6c7843f2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_saturate.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""paged_cache.fp8_e4m3_saturate (FP8_SATURATE): the E4M3 cache writes and the transient FP8 query casts saturate to ++-448 like the CUDA writers (__NV_SATFINITE); every finite in-range value keeps the plain cast's bits, values already +in E4M3 pass through, and the kill switch is exactly the plain cast. Where torch's cast already saturates (this image) +the conversion is the cast alone, one kernel; otherwise it clamps first, the same bytes.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import paged_cache + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fp8_e4m3_saturate_matches_satfinite(dtype, monkeypatch) -> None: + g = torch.Generator(device="cuda").manual_seed(0) + inside = (torch.randn(4096, generator=g, device="cuda") * 100).clamp(-448, 448).to(dtype) + plain = inside.to(torch.float8_e4m3fn) + assert torch.equal( + paged_cache.fp8_e4m3_saturate(inside).view(torch.uint8), plain.view(torch.uint8) + ) + edge = torch.tensor( + [449.0, 463.0, 464.0, 1e4, float("inf"), -500.0, -float("inf")], dtype=dtype, device="cuda" + ) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0], device="cuda") + assert torch.equal(paged_cache.fp8_e4m3_saturate(edge).float(), want) + nan = paged_cache.fp8_e4m3_saturate(torch.tensor([float("nan")], dtype=dtype, device="cuda")) + assert torch.isnan(nan.float()).all() + already = edge.clamp(-448, 448).to(torch.float8_e4m3fn) + assert paged_cache.fp8_e4m3_saturate(already) is already + monkeypatch.setattr(paged_cache, "FP8_SATURATE", False) + killed = paged_cache.fp8_e4m3_saturate(edge) + assert torch.equal(killed.view(torch.uint8), edge.to(torch.float8_e4m3fn).view(torch.uint8)) + + +def test_cache_write_saturates() -> None: + """The cache-dtype conversion every write_kv_slots branch applies (compared as bytes: torch has no float8 + index_copy_ kernel), and an E4M3 source passing through untouched.""" + values = torch.tensor([[1.0, 500.0, -1e5, float("inf")]], dtype=torch.bfloat16, device="cuda") + out = paged_cache._to_cache_dtype(values, torch.float8_e4m3fn) + assert out.dtype == torch.float8_e4m3fn + want = torch.tensor([[1.0, 448.0, -448.0, 448.0]], device="cuda").to(torch.float8_e4m3fn) + assert torch.equal(out.view(torch.uint8), want.view(torch.uint8)) + assert paged_cache._to_cache_dtype(out, torch.float8_e4m3fn) is out + bf = paged_cache._to_cache_dtype(values, torch.bfloat16) + assert bf.dtype == torch.bfloat16 and torch.equal(bf, values) + + +_EDGES = [ + 448.0, + 449.0, + 463.0, + 464.0, + 1e4, + float("inf"), + -448.0, + -449.0, + -464.0, + -1e4, + -float("inf"), + float("nan"), + 0.0, +] + + +def _cuda_kernels(fn) -> int: + from torch.profiler import ProfilerActivity, profile + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + return sum(1 for e in prof.events() if e.device_type.name == "CUDA") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("cast_saturates", [True, False], ids=["cast_only", "clamp_fallback"]) +def test_fp8_e4m3_saturate_bytes_equal_clamp_then_cast(dtype, cast_saturates, monkeypatch) -> None: + """Both paths give the bytes of clamp(+-448) followed by the cast, edge values included.""" + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", cast_saturates) + x = torch.tensor(_EDGES, dtype=dtype, device="cuda") + want = x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + assert torch.equal(paged_cache.fp8_e4m3_saturate(x).view(torch.uint8), want.view(torch.uint8)) + + +def test_fp8_e4m3_saturate_takes_the_cast_only_path_here(monkeypatch) -> None: + """This image's torch cast saturates, on the CPU (the import-time probe) and on CUDA, so the conversion is one + kernel; with the probe failing it falls back to clamp + cast, two kernels.""" + assert paged_cache._CAST_SATURATES + x = torch.tensor(_EDGES, dtype=torch.bfloat16, device="cuda") + assert torch.equal( + x.to(torch.float8_e4m3fn).view(torch.uint8), + x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).view(torch.uint8), + ) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 1 + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", False) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 2 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py new file mode 100644 index 000000000000..859e28a464d8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for the attention/m3_fullctx_decode entry: the dense layers' full-context FP8 decode on PrimTS +split-KV decode at the recipe the target pins per (query rows per request, requests). + +Every cell runs against the fp32 reference (q rounded to E4M3 first, as the kernel runs it), eager and replayed from a +CUDA graph (the same bits), and twice (the same bits). The K/V pages are read through the strided views of a flat +sub-page pool, as the KV cache manager hands them out. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_fullctx_decode import ( # noqa: E501 + HEAD_DIM, + HEADS, + KV_HEADS, + MAX_KV_LEN, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_fullctx_decode import ( # noqa: E501 + m3_fullctx_decode as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) + +PAGE = 128 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the entry asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor, subpages: int): + self.pool = pool + self.subpages = subpages + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, self.subpages + + +def _inputs(kv_lens, qlen: int, seed: int, subpages: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * subpages, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1, so a split that missed the tail fails the band. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + for page in range(start // PAGE, (n + PAGE - 1) // PAGE): + v = int(block_table[b, page]) * subpages + 1 + pool[v, :, (start % PAGE if page == start // PAGE else 0) :] += 1.0 + # q as the layer hands it over on decode steps: the FP8 query columns of the fused norm/RoPE kernel's + # [q | k | v] output (row stride wider than q), so no cast makes it compact. + qkv = torch.randn( + (batch * qlen, (HEADS + 2 * KV_HEADS) * HEAD_DIM), device="cuda", generator=g + ).to(torch.float8_e4m3fn) + q = qkv[:, : HEADS * HEAD_DIM].view(batch * qlen, HEADS, HEAD_DIM) + seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32) + return q, _Pool(pool.to(torch.float8_e4m3fn), subpages), block_table, seq_lens + + +@pytest.mark.parametrize("subpages", [2, 3], ids=["kv", "kv+index"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, subpages, certifying) -> None: # noqa: F811 + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], qlen, zlib.crc32(cell.why.encode()) % 2**31, subpages + ) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + certifying(op) + ref = op.reference(*args) + eager = op(*args).clone() + torch.testing.assert_close(eager.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + again = op(*args).clone() + assert torch.equal(again, eager), "not deterministic" + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args) + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager), "the CUDA-graph replay differs from the eager call" + + +def test_the_cells_cover_the_target_rows() -> None: + certified = {(c.spec["qlen"], len(c.spec["kv_lens"])) for c in op.CELLS} + assert set(modeling.FULLCTX_DECODE.keys) <= certified + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: + """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once + in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + + if not kernel.has_own_reducer(): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_fullctx_reduce") + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs(spec["kv_lens"], qlen, 11, 2) + outs = [] + for own in (True, False): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + op(q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe, own_reducer=own) + outs.append(out.float()) + ours, primts = outs + _, exp = torch.frexp(primts.abs()) + spacing = torch.where( + primts == 0, + torch.full_like(primts, 2.0**-133), + torch.ldexp(torch.ones_like(primts), exp - 8), + ) + assert bool(((ours - primts).abs() <= spacing).all()) + + +def _tile_edge_kv_lens(plan, qlen: int): + """KV lengths around every kind of boundary the split ranges have, so each query row of the Q CTA in turn ends + its keys on a KV tile edge (k * tile + d, d in -1 .. qlen - 1): the first tiles, the reducer's 8-group stride, the + lengths where a split's range grows (splits * insts tiles) and a split edge past it.""" + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: + """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys + (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode + step first hit (NaN rows from exp2(-inf - -inf)).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, False, get_env_enable_pdl() + ) + lens = _tile_edge_kv_lens(kernel.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, sum(kv_lens), 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: + """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just + past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + lengths = [(1 << 20) + 1, (1 << 20) + 8] + ([1_500_001] if nreq == 1 else []) + batches = [[n] * nreq if nreq == 1 else [n, (1 << 20) - 3] for n in lengths] + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, 13, 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, own, get_env_enable_pdl() + ) + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +@pytest.mark.parametrize( + "splits,rows,seq_len_q", [(108, 8, 1), (64, 8, 1), (96, 32, 4), (96, 32, 4)] +) +def test_the_tp8_merge_is_the_tp2_merge_bitwise(splits, rows, seq_len_q) -> None: + """The TP8 copy of the BF16 split-KV merge (one row per CTA, early dependent release) returns the image TP2 op's + bits on random 16-bit partials with empty (-inf log2-LSE) splits, both partial dtypes, two requests.""" + ops = torch.ops.trtllm + if not ( + hasattr(ops, "minimax_m3_sm107tp8_fullctx_reduce") + and hasattr(ops, "minimax_m3_sm107_tp2_fullctx_reduce") + ): + pytest.skip("needs both the TP8 merge copy and the image's TP2 merge op") + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(splits * 131 + rows) + nreq, kvh = 2, 1 + for dtype in (torch.float16, torch.bfloat16): + po = ( + (torch.randn(nreq, kvh, splits, rows, 128, device=dev, generator=g) * 3) + .to(dtype) + .contiguous() + ) + ps = torch.randn(nreq, kvh, splits, rows, device=dev, generator=g) * 4 + ps[torch.rand(ps.shape, device=dev, generator=g) < 0.05] = float("-inf") + seq = torch.tensor([165521, 120011], dtype=torch.int32, device=dev) + outs = [] + for reduce in ( + ops.minimax_m3_sm107_tp2_fullctx_reduce, + ops.minimax_m3_sm107tp8_fullctx_reduce, + ): + out = torch.zeros( + nreq * seq_len_q, kvh * rows // seq_len_q, 128, dtype=torch.bfloat16, device=dev + ) + reduce( + po, ps.contiguous(), seq, out, seq_len_q, seq_len_q, 128, 1, seq_len_q > 1, False + ) + outs.append(out) + torch.cuda.synchronize() + assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)), ( + splits, + rows, + seq_len_q, + dtype, + ) + + +def _swizzled_scale_offsets(rows: int, width: int) -> torch.Tensor: + """Offsets of rows [0, rows)'s UE8M0 scales in the 128x4-swizzled layout of mxfp8_quantize(x, True).""" + k_sf = width // 32 + k_tiles = (k_sf + 3) // 4 + r = torch.arange(rows, device="cuda").view(-1, 1) + k = torch.arange(k_sf, device="cuda").view(1, -1) + offsets = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + return offsets.reshape(-1) + + +@pytest.mark.parametrize("splits,rows,seq_len_q", [(108, 8, 1), (64, 8, 1), (96, 32, 4)]) +def test_the_tp8_merge_mxfp8_copy_is_stock_mxfp8_quantize(splits, rows, seq_len_q) -> None: + """The TP8 merge's MXFP8 copy (the dense o_proj hand-off) equals mxfp8_quantize(out, True) bit for bit: E4M3 rows + and the real rows' swizzled UE8M0 scales, with an all-zero head row (amax 0) and saturating outliers; the BF16 + output is the merge without the copy.""" + ops = torch.ops.trtllm + if not hasattr(ops, "minimax_m3_sm107tp8_fullctx_reduce"): + pytest.skip("needs the TP8 merge copy") + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(splits * 7 + rows) + nreq, kvh = 2, 1 + tokens = nreq * seq_len_q + width = kvh * rows // seq_len_q * 128 + for dtype in (torch.float16, torch.bfloat16): + po = torch.randn(nreq, kvh, splits, rows, 128, device=dev, generator=g) * 3 + po[0, 0, :, 1] = 0.0 + po[1, 0, :, 0, 5] = 6.0e4 + po = po.to(dtype).contiguous() + ps = (torch.randn(nreq, kvh, splits, rows, device=dev, generator=g) * 4).contiguous() + seq = torch.tensor([165521, 120011], dtype=torch.int32, device=dev) + args = (po, ps, seq) + tail = (seq_len_q, seq_len_q, 128, 1, seq_len_q > 1, False) + plain = torch.zeros(tokens, width, dtype=torch.bfloat16, device=dev) + ops.minimax_m3_sm107tp8_fullctx_reduce(*args, plain, *tail) + out = torch.zeros(tokens, width, dtype=torch.bfloat16, device=dev) + q8 = torch.empty(128, width, dtype=torch.float8_e4m3fn, device=dev) + sf = torch.zeros(128 * (width // 32), dtype=torch.uint8, device=dev) + ops.minimax_m3_sm107tp8_fullctx_reduce(*args, out, *tail, q8, sf) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), plain.view(torch.int16)) + stock_q, stock_sf = ops.mxfp8_quantize(out, True) + assert torch.equal( + q8[:tokens].view(torch.uint8), stock_q.view(torch.uint8).reshape(tokens, -1) + ), dtype + at = _swizzled_scale_offsets(tokens, width) + assert torch.equal(sf[at], stock_sf.view(torch.uint8).reshape(-1)[at]), dtype diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_identity_selection.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_identity_selection.py new file mode 100644 index 000000000000..8819b1fcfb73 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_identity_selection.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the short-sequence identity selection (runtime policy ``short_sequence_identity``) is the stock table. + +On a step whose every query token has at most 16 valid blocks, this target skips the index scorer and the select and +stages the table (i < n_valid ? i : -1). Stock runs the in-tree CuTe DSL scorer and ``select_blocks_from_maxscore``, +which returns ascending ids and drops a valid block only when its max score is -inf. The scorer's scores are FP8 +products summed in FP32 (|score| <= 448^2 * 128), so -inf comes only from the causal mask, which n_valid excludes. +This test drives the real scorer with extreme FP8 operands at the TP8 geometry (1 index head per rank) and checks +both claims. +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) +from tensorrt_llm._torch.custom_ops import ( + cute_dsl_custom_ops, # noqa: F401 — registers the scorer op +) + +PAGE, DIM, TOPK, HEADS, QLEN = 128, 128, 16, 1, 4 +E4M3_MAX = 448.0 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + + +def _fill(kind, shape, gen): + if kind == "random": + x = torch.randn(shape, generator=gen, device="cuda") * 64 + elif kind == "max": + x = torch.full(shape, E4M3_MAX, device="cuda") + elif kind == "negmax": + x = torch.full(shape, -E4M3_MAX, device="cuda") + elif kind == "alternating": + x = torch.full(shape, E4M3_MAX, device="cuda") + x[..., 1::2] = -E4M3_MAX + elif kind == "zero": + x = torch.zeros(shape, device="cuda") + else: # the smallest E4M3 subnormal + x = torch.full(shape, 2.0**-9, device="cuda") + return x.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + + +_KINDS = ("random", "max", "negmax", "alternating", "zero", "tiny") + + +@pytest.mark.parametrize( + "seq_lens", [(4,), (129,), (2048,), (700, 2048), (5, 1500, 2047, 300), (1, 128, 256, 1024)] +) +@pytest.mark.parametrize("q_kind", _KINDS) +@pytest.mark.parametrize("k_kind", _KINDS) +def test_stock_selection_of_a_short_step_is_the_identity_table(seq_lens, q_kind, k_kind) -> None: + gen = torch.Generator(device="cuda").manual_seed(0) + reqs, total_q = len(seq_lens), len(seq_lens) * QLEN + nblk = [(s + PAGE - 1) // PAGE for s in seq_lens] + pages = sum(nblk) + 4 + perm = torch.randperm(pages, generator=gen, device="cuda").int() + block_table = torch.zeros((reqs, max(nblk)), dtype=torch.int32, device="cuda") + start = 0 + for r, n in enumerate(nblk): + block_table[r, :n] = perm[start : start + n] + start += n + n_valid = torch.tensor( + [(max(s - QLEN + o + 1, 0) + PAGE - 1) // PAGE for s in seq_lens for o in range(QLEN)], + dtype=torch.int32, + device="cuda", + ) + idx_q = _fill(q_kind, (total_q, HEADS, DIM), gen) + idx_k = _fill(k_kind, (pages, PAGE, DIM), gen) + max_score = torch.full((HEADS, max(nblk), total_q), float("-inf"), device="cuda") + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, + idx_k, + block_table, + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + max_score.transpose(1, 2), + QLEN, + ) + for tok, nv in enumerate(n_valid.tolist()): + assert bool(torch.isfinite(max_score[:, :nv, tok]).all()), ( + f"token {tok}: a causally valid block scored non-finite" + ) + table = select_blocks_from_maxscore( + max_score, + topk=TOPK, + n_valid_blocks=n_valid, + init_blocks=1, + local_blocks=4, + head_major_output=True, + ).to(torch.int32) + # The selector returns [tokens, kv_heads, topk]; the decode reads it as [kv_heads, tokens, topk] (layers.py). + table = table.permute(1, 0, 2) + i = torch.arange(TOPK, dtype=torch.int32, device="cuda") + identity = torch.where( + i[None, :] < n_valid[:, None], i[None, :], torch.full_like(i[None, :], -1) + ) + assert torch.equal(table, identity[None].expand(HEADS, total_q, TOPK)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..e768b7442a15 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score entry: the MSA index scorer with its pre-wait schedule. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score`` at its own 256-CTA split) bit for bit, which the +pre-wait prefetch, the early tile issue and the cell's split must not change. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_index_score import ( # noqa: E501 + PAGE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 1 # index heads per rank at TP8 (4 over 4 KV groups, each group on a rank pair) +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + assert spec["heads"] == HEADS + args = (idx_q, index_k, block_table, seq, score, qlen, spec["split_k"], spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, 1536, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4 + ) + assert torch.equal(score, stock) + + +@pytest.mark.parametrize("qlen", [1, 2, 3, 4]) +def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> None: + """Lengths around the edge between pre-staged (all-old) and post-wait first blocks, 1 and 2 requests.""" + for k in (1, 2, 1294): + for length in (k * PAGE - qlen - 1, k * PAGE - qlen, k * PAGE - qlen + 1, k * PAGE + 1): + for seq_lens, split in ( + ((length,), 1536), + ((length,), 1272), + ((length, max(qlen, length // 3)), 768), + ): + idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) + score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) + op(idx_q, index_k, block_table, seq, score, 4, split, True) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + assert torch.equal(score, stock), (qlen, seq_lens) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, 1536, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..0539d9639125 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_msa_attention entry: one layer's MSA backend step against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp8``): a hybrid NVFP4 ``MiniMaxM3KVCacheManagerV2`` +with the requests allocated in it, ``MiniMaxM3MsaSparseAttentionMetadata`` filled by its ``prepare()``, and the layer's +``MiniMaxM3MsaSparseAttention`` built as the target builds it. Earlier rows (and, on a sparse layer, the step's own +rows, which the target's producer writes before the call) go into the cache through the in-tree scatters, at the +slots the manager assigned. The sparse layers' index-K rows are planted so every block's score is 1/8 or more apart +from every other's: the top-k selection the reference recomputes is then unambiguous. + +The cells are the target's calls (``_attention_core`` in modeling.py): a dense layer with live bf16 K / V in a +prefill chunk and live FP8 K / V in a decode step (the dense layers sit in the hybrid cache's FP8 half), a sparse +layer whose producer already wrote the caches (``k = v = None``) in a prefill chunk beyond 16 blocks and in a mixed +step, both with the NVFP4 scales the target passes. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_msa_attention import ( # noqa: E501 + m3_msa_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp8 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" + +# Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, +# per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 +# FMHA does), so what is left is each FMHA's own arithmetic, measured on GB300 over these cells (max element / RMS): +# * fmha_sm100 over the FP8 pool (dense context rows): 10.7 / 6.0. It runs E4M3 x E4M3 and its softmax +# probabilities are rounded before the PV product. +# * trtllm-gen over the FP8 pool (dense generation rows): 7.7 / 6.1. It rounds the probabilities to E4M3 (scaled by +# 448 against the running maximum): emulating that brings a single-tile request to 0.6 RMS. +# * the NVFP4 CSR prefill (sparse context rows): 33.6 / 11.8 at 4096 rows, 24.3 / 9.1 at 512. It stages the +# reconstructed K / V in E4M3 (the in-tree A/B test, test_minimax_m3_sparse_attn_decode, notes it and holds it +# to 3e-2). +# * the Triton NVFP4 sparse decode (sparse generation rows): 1.0 / 0.19 -- the band of its own entry +# (m3_sparse_decode): exact operands, FP16 probabilities, BF16 rounding. +# Each band leaves about 1.5x the measured maximum; a shifted or wrongly selected reference is far outside them. +_FP8_CONTEXT = dict(element_ulp=16.0, rms_ulp=9.0) +_FP8_DECODE = dict(element_ulp=12.0, rms_ulp=9.0) +_NVFP4_CONTEXT = dict(element_ulp=48.0, rms_ulp=16.0) +_NVFP4_DECODE = dict(element_ulp=2.0, rms_ulp=0.5) +_BANDS = { + ("dense", "context"): (_FP8_CONTEXT, None), + ("dense", "generation"): (None, _FP8_DECODE), + ("sparse", "context"): (_NVFP4_CONTEXT, None), + ("sparse", "mixed"): (_NVFP4_CONTEXT, _NVFP4_DECODE), +} + + +def _assert_in_bands(output, ref, bands, context_rows: int) -> None: + """Gate the context rows and the generation rows each on their own kernel's band.""" + context_band, decode_band = bands + if context_rows: + assert_within_ulp(output[:context_rows], ref[:context_rows], **context_band) + if context_rows < output.shape[0]: + assert_within_ulp(output[context_rows:], ref[context_rows:], **decode_band) + + +# (request rows, cached tokens per request, new tokens per request, context rows) per cell. +_SHAPES = { + ("dense", "context"): ((384,), (1024,), 1), + ("dense", "generation"): ((128, 2500), (1, 1), 0), + ("sparse", "context"): ((0,), (4096,), 1), + ("sparse", "mixed"): ((0, 2999, 4500), (512, 1, 1), 1), +} + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """The two quantities ``assert_within_ulp`` gates, in BF16 ulps.""" + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _shift(x: torch.Tensor) -> torch.Tensor: + """``x``'s rows moved down by one (byte-wise, so E4M3 rows move too).""" + return x.view(torch.uint8).roll(1, 0).view(x.dtype) if x.element_size() == 1 else x.roll(1, 0) + + +def _shape(cell): + spec = cell.spec + cached, qlens, contexts = _SHAPES[spec["layer"], spec["phase"]] + if spec["phase"] == "context": + assert sum(qlens) == spec["tokens"] + elif spec["phase"] == "generation": + assert len(qlens) == spec["reqs"] + else: + assert sum(qlens[:contexts]) == spec["tokens"] and len(qlens) - contexts == spec["reqs"] + return cached, qlens, contexts + + +def _run_dense(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + rt.write_fp8_rows(manager, rt.DENSE, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + # A prefill chunk arrives in BF16 (the hybrid cache's dense layers take FP8 Q / K / V only on decode steps). + dtype = torch.bfloat16 if contexts else torch.float8_e4m3fn + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + k = torch.randn(total, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + v = torch.randn(total, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + # The NVFP4 target passes its layer scales on every layer; an FP8 layer's write and FMHA take none. + inv = torch.tensor([1.0, 0.5, 2.0], device="cuda") + fwd = 1.0 / inv + return (q, k, v, None, inv, fwd), dict(rids=rids, cached=cached, qlens=qlens, contexts=contexts) + + +def _run_sparse(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + u0, u1 = rt.index_directions(gen) + lengths = [c + n for c, n in zip(cached, qlens)] + keys = [ + torch.randn(n, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths + ] + values = [ + torch.randn(n, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths + ] + quant = rt.nvfp4_quant(torch.cat(keys), torch.cat(values)) + for rid, n, kk, vv in zip(rids, lengths, keys, values): + # Every row, the step's own included: the producer writes a sparse layer's caches before the call. + idx_k = rt.planted_index_keys(n, u0, u1, gen) + rt.write_nvfp4_rows( + manager, + metadata, + rt.SPARSE, + rid, + torch.arange(n), + kk.to(torch.bfloat16), + vv.to(torch.bfloat16), + idx_k, + quant, + ) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + # Query tokens alternate between the two planted directions, so neighbouring tokens select different blocks; the + # rank's index heads take alternating directions, so with more than one KV head they select different blocks too. + turn = torch.arange(total, device="cuda") % 2 + directions = torch.stack([u0, u1]) + idx_q = torch.cat([directions[(turn + h) % 2] for h in range(rt.INDEX_HEADS)], dim=1) + idx_q = idx_q.to(torch.float8_e4m3fn) + return (q, None, None, idx_q, quant, 1.0 / quant), dict( + rids=rids, cached=cached, qlens=qlens, contexts=contexts + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + sparse = spec["layer"] == "sparse" + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=False) + try: + backend = rt.make_msa_backend(rt.SPARSE if sparse else rt.DENSE) + metadata = rt.make_msa_metadata(manager) + run = _run_sparse if sparse else _run_dense + (q, k, v, idx_q, inv, fwd), step = run(cell, gen, manager, backend, metadata) + output = torch.full( + (q.shape[0], rt.HEADS * rt.HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + certifying(op) + result = op(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert result is output + ref = op.reference(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert torch.isfinite(output.float()).all(), "every live row is written" + bands = _BANDS[spec["layer"], spec["phase"]] + context_rows = sum(step["qlens"][: step["contexts"]]) + for name, rows in ( + ("context", slice(0, context_rows)), + ("generation", slice(context_rows, None)), + ): + if output[rows].numel(): + element, rms = _ulp_stats(output[rows], ref[rows]) + print( + f"\n[{cell.why}] {name} rows: max element {element:.2f} ulp, rms {rms:.3f} ulp" + ) + _assert_in_bands(output, ref, bands, context_rows) + + if k is not None: + # The call's own cache write: the live K / V rows land, as a plain E4M3 cast, at their positions. + pool = manager.get_buffers(rt.DENSE, "HND") + first = 0 + for rid, c, n in zip(step["rids"], step["cached"], step["qlens"]): + slots = rt.flat_slots(manager, rt.DENSE, rid, torch.arange(c, c + n)).long() + for role, src in ((0, k), (1, v)): + for h in range(rt.KV_HEADS): + got = pool[slots // rt.PAGE, role, h, slots % rt.PAGE] + want = src[first : first + n, h * rt.HEAD_DIM : (h + 1) * rt.HEAD_DIM] + want = want.to(torch.float8_e4m3fn) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose query rows are shifted by one token is outside it. + shifted = op.reference(backend, _shift(q), k, v, idx_q, metadata, output, sparse, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, shifted, bands, context_rows) + if sparse: + # And the selection is load-bearing: attending every causal block instead is outside it too. + dense = op.reference(backend, q, k, v, idx_q, metadata, output, False, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, dense, bands, context_rows) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py new file mode 100644 index 000000000000..baf8ef7a2b73 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-token valid-block count is written only before the step's first layer. + +The long-row select (``kernels/m3_select_long_dsl.py``, ``SELECT_N_VALID_PREWAIT``) reads its row's entry of +``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the +buffer happens before the step's first layer: + +- the allocation (``_alloc_msa_proxy_scratch``); +- the host-to-device copy in ``prepare()`` (``_build_step_plans``); +- ``on_update_kv_lens``, which the runner calls before the forward. It writes through the plain + ``<<<>>>`` launch of ``msa_kv_update_kernel`` or through torch copies, so any PDL grid of the step starts only + after it completes. + +This test scans ``tensorrt_llm/_torch`` and fails on any other write site (including writes through a local alias of +the buffer), so a new writer has to revisit the pre-wait read. A buffer handed to an opaque kernel launch is not +tracked; such a launch inside a layer must be reviewed by hand. + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +_NAME = "msa_n_valid_blocks" +# In-place tensor methods that write their receiver. +_MUTATORS = { + "copy_", + "fill_", + "zero_", + "index_put_", + "index_copy_", + "scatter_", + "masked_fill_", + "add_", + "sub_", +} +# Calls whose first argument is the destination they write: ``step_inputs.copy_(dst, src)`` and the H2D helpers. +_COPY_HELPERS = {"copy_", "copy_h2d_if_changed", "copy_h2d_pool_shared"} +_BACKEND = "attention/backends/sparse/minimax_m3/msa_backend.py" +# (file relative to tensorrt_llm/_torch, enclosing function) -> writes allowed there +_ALLOWED = { + (_BACKEND, "_alloc_msa_proxy_scratch"), # the allocation + (_BACKEND, "_build_step_plans"), # the H2D copy in prepare() + (_BACKEND, "on_update_kv_lens"), # the device-side correction before the forward +} + + +def _base(node): + """The innermost receiver of an attribute / subscript / call chain.""" + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute) and node.attr == _NAME: + return node + node = node.func if isinstance(node, ast.Call) else node.value + return node + + +def _mentions(node, aliases) -> bool: + """Whether the expression's receiver chain is ``*.msa_n_valid_blocks`` or a local alias of it.""" + b = _base(node) + if isinstance(b, ast.Attribute): + return b.attr == _NAME + return isinstance(b, ast.Name) and (b.id == _NAME or b.id in aliases) + + +def _refers(node, aliases) -> bool: + """Whether any sub-expression (e.g. a branch of a conditional) names the buffer or an alias.""" + return any( + (isinstance(n, ast.Attribute) and n.attr == _NAME) + or (isinstance(n, ast.Name) and (n.id == _NAME or n.id in aliases)) + for n in ast.walk(node) + ) + + +def _write_sites(tree: ast.AST): + sites = [] + + def visit(node, func, aliases): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + func, aliases = node.name, set() + if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + value = node.value + for t in targets: + if isinstance(t, ast.Name) and value is not None and _refers(value, aliases): + aliases.add(t.id) # a local view of the buffer + elif ( + not isinstance(t, ast.Name) + and _mentions(t, aliases) + and not isinstance(node, ast.AnnAssign) + ): + sites.append((func, node.lineno)) + elif ( + isinstance(t, ast.Attribute) + and t.attr == _NAME + and isinstance(node, ast.AnnAssign) + ): + pass # a dataclass field declaration + if isinstance(node, ast.Call): + f = node.func + if isinstance(f, ast.Attribute) and f.attr in _MUTATORS and _mentions(f.value, aliases): + sites.append((func, node.lineno)) + name = ( + f.id + if isinstance(f, ast.Name) + else (f.attr if isinstance(f, ast.Attribute) else "") + ) + if name in _COPY_HELPERS and node.args and _mentions(node.args[0], aliases): + sites.append((func, node.lineno)) + for child in ast.iter_child_nodes(node): + visit(child, func, aliases) + + visit(tree, None, set()) + return sites + + +def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: + found = set() + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if _NAME not in src: + continue + rel = path.relative_to(_ROOT).as_posix() + for func, line in _write_sites(ast.parse(src)): + assert (rel, func) in _ALLOWED, ( + f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " + "wait (kernels/m3_select_long_dsl.py, SELECT_N_VALID_PREWAIT), so a writer inside the " + "step must be ordered before the select's grid or the read moved after the wait" + ) + found.add((rel, func)) + assert found == _ALLOWED, f"expected write sites {_ALLOWED}, found {found}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..124942104b82 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_producer entry: the pre-wait-staged NVFP4 QKV + indexer producer. + +Every cell is held to the native producer (``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert``) bit for bit on q, +index-q and every K / V / scale / index-K cache byte, and to the fp32 torch reference: the FP8 outputs within one E4M3 +step, the decoded NVFP4 K / V within one E2M1 step of their block scale. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_producer import ( # noqa: E501 + FP4_BLOCK, + HEADS_Q, + PAGE, + decode_nvfp4_token, + scale_index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_producer import ( # noqa: E501 + m3_producer as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_producer requires a CUDA device" + +PAGES = 1024 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _caches(): + return ( + torch.zeros(PAGES, 2, 1, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, 1, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _inputs(tokens: int, scale: float, bad_slots: bool, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + weights = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + packed = ( + torch.randn(tokens, (HEADS_Q + 4) * 128, device="cuda", generator=g) * scale + ).bfloat16() + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return packed, slots, quant, weights, rotary, positions + + +def _args(inputs, caches, early: bool = True): + packed, slots, quant, weights, rotary, positions = inputs + return (packed, *caches, slots, quant, EPS, *weights, rotary, positions, HEADS_Q, True, early) + + +def _native(inputs, caches): + packed, slots, quant, weights, rotary, positions = inputs + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, *caches, slots, quant, HEADS_Q, 1, 1, 128, 64, EPS, *weights, rotary, positions + ) + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + early = spec.get("early", True) + inputs = _inputs( + spec["tokens"], spec["scale"], spec.get("bad_slots", False), zlib.crc32(cell.why.encode()) + ) + stock_caches, caches = _caches(), _caches() + stock = _native(inputs, stock_caches) + args = _args(inputs, caches, early) + certifying(op) + ref = op.reference(*args) + out = op(*args) + torch.cuda.synchronize() + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, s, o in zip(names, (*stock, *stock_caches), (*out, *caches)): + assert torch.equal(s.view(torch.uint8), o.view(torch.uint8)), ( + f"{name} differs from the native producer" + ) + + _within_fp8(out[0], ref["q"], "q") + _within_fp8(out[1][:, 0], ref["index_q"], "index_q") + kv_data, kv_scale, index_k = caches + slots = inputs[1] + checked = 0 + for t in range(spec["tokens"]): + slot = int(slots[t]) + if slot < 0 or slot >= PAGES * PAGE: + continue + page, within = divmod(slot, PAGE) + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") + for plane, name in ((0, "k"), (1, "v")): + scales = kv_scale[page, plane, 0].reshape(-1)[scale_index(within, plane).cuda()] + dq = decode_nvfp4_token(kv_data[page, plane, 0, within], scales, KV_QUANT[1 + plane]) + r = ref[name][t].bfloat16().float() + # Round to nearest E2M1 under the block's E4M3-rounded scale s (in value units s / quant_scale): at most + # half the widest E2M1 step (one unit), plus the clipping at 6 units when s rounded down; a scale that + # underflows E4M3 zeroes the block. A slack of one E4M3 step covers the kernel's approximate reciprocals. + qs = KV_QUANT[1 + plane] + amax = r.view(-1, FP4_BLOCK).abs().amax(dim=1) + unit = (qs * amax / 6.0).to(torch.float8_e4m3fn).float() / qs + bound = torch.where( + unit > 0, 1.15 * unit + (amax - 6.0 * unit / 1.15).clamp_min(0.0), amax + ) + bound = bound.repeat_interleave(FP4_BLOCK) + 1e-12 + assert bool(((dq - r).abs() <= bound).all()), (name, t) + checked += 1 + assert checked >= spec["tokens"] - 2 + + +def test_early_trigger_does_not_change_the_bits() -> None: + inputs = _inputs(48, 1.0, False, 7) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=True)) + out_b = op(*_args(inputs, b, early=False)) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +@pytest.mark.parametrize( + "tokens, scale, bad_slots, early", + [ + (1, 1.0, False, True), + (4, 3.0, False, True), + (16, 1e-3, False, True), + (32, 40.0, True, True), + (48, 1.0, False, True), + (1000, 1.0, False, False), + ], +) +def test_prewait_step_inputs_is_bitwise_the_shared_op(tokens, scale, bad_slots, early) -> None: + """modeling's PRODUCER_PREWAIT_STEP_INPUTS: this target's copy (minimax_m3_sm107_tp8_producer_pw, the position ids, + slots and RoPE rows read before its grid wait) writes exactly the shared op's q, index-q and cache bytes.""" + assert hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_producer_pw"), ( + "the target's producer op is not loaded (its TP8 extension)" + ) + inputs = _inputs(tokens, scale, bad_slots, 11) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=early)) + out_b = op(*_args(inputs, b, early=early), prewait_step_inputs=True) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +def test_guard_refuses_a_repeated_slot() -> None: + inputs = list(_inputs(4, 1.0, False, 3)) + inputs[1] = inputs[1].clone() + inputs[1][1] = inputs[1][0] + with pytest.raises(AssertionError, match="must not repeat a slot"): + with validating(op): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_operand_off_the_device() -> None: + inputs = list(_inputs(4, 1.0, False, 9)) + inputs[5] = inputs[5].cpu() # position ids on the host + with pytest.raises(RuntimeError, match="packed's CUDA device"): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_unaligned_index_k_page_stride() -> None: + """The index-K rows are written as 32-bit words: a page stride that is not a multiple of 4 is refused.""" + kv_data, kv_scale, _ = _caches() + backing = torch.zeros(PAGES * (128 * 128 + 1), dtype=torch.float8_e4m3fn, device="cuda") + index_k = backing.as_strided((PAGES, 1, PAGE, 128), (128 * 128 + 1, 128 * 128, 128, 1)) + with pytest.raises(RuntimeError, match="index_k"): + op(*_args(_inputs(4, 1.0, False, 13), (kv_data, kv_scale, index_k))) + + +def test_op_refuses_unaligned_weights_or_rope_table() -> None: + """The norm weights are read as 8-byte words and the RoPE rows as float4: offset views are refused.""" + packed, slots, quant, weights, rotary, positions = _inputs(4, 1.0, False, 17) + offset_weight = torch.zeros(129, dtype=torch.bfloat16, device="cuda")[1:] + with pytest.raises(RuntimeError, match="norm weights"): + op( + *_args( + (packed, slots, quant, [offset_weight, *weights[1:]], rotary, positions), _caches() + ) + ) + flat = torch.zeros(rotary.numel() + 1, device="cuda") + offset_rope = flat[1:].view(rotary.shape) + with pytest.raises(RuntimeError, match="rotary_cos_sin"): + op(*_args((packed, slots, quant, weights, offset_rope, positions), _caches())) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..6e58bacee51c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_qk_norm_rope_fp8 entry of the MiniMax-M3 tp8 target: the out-of-place FP8 form of the fused +per-head Gemma RMSNorm + partial NeoX RoPE against the in-place kernel followed by the E4M3 cast. V is a cast, so it +matches bit for bit; Q and K round once from fp32 where the composition rounds through bf16 first, so each element is +within one E4M3 step of it.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 + m3_qk_norm_rope_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" + +HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 + + +def _args(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q, kv = spec["q"], spec["kv"] + qkv = torch.randn(spec["m"], (q + 2 * kv) * HEAD_DIM, generator=g, device="cuda").to( + torch.bfloat16 + ) + qw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (spec["m"],), generator=g, device="cuda", dtype=torch.int32) + return ( + qkv, + q, + kv, + kv, + HEAD_DIM, + ROTARY_DIM, + EPS, + qw, + kw, + THETA, + True, + pos, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = args[0].clone() + with validating(op): + out = op(*args) + assert torch.equal(args[0], before), "the FP8 form must leave qkv untouched" + ref = op.reference(*args) + assert out.dtype == torch.float8_e4m3fn and out.shape == ref.shape + v_start = (cell.spec["q"] + cell.spec["kv"]) * HEAD_DIM + assert torch.equal(out[:, v_start:].view(torch.uint8), ref[:, v_start:].view(torch.uint8)) + assert _e4m3_steps(out[:, :v_start], ref[:, :v_start]) <= 1 + + +def _e4m3_steps(a, b) -> int: + """The largest distance between two E4M3 tensors in representable values (E4M3 codes order by magnitude).""" + + def signed(t): + c = t.view(torch.uint8).int() + # Sign-magnitude to a line: -0 and +0 both map to 0. + return torch.where(c >= 128, 128 - c, c) + + return int((signed(a) - signed(b)).abs().max()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..514567365f8c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_qkv_insert entry: the generic horizontal producer on an FP8 main cache. + +Q, index-Q and every written K / V / index-K cache row are held within one E4M3 step of the fp32 reference. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_qkv_insert import ( # noqa: E501 + m3_qkv_insert as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" + +PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 8, 1, 1, 128, 64 +MAX_POSITIONS = 262144 +EPS = 1e-6 + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +def _args(tokens: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + tokens, (HEADS + 2 * KV + INDEX + 1) * D, device="cuda", generator=g + ).bfloat16() + kv = torch.zeros(PAGES, 2, KV, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + index_k = torch.zeros(PAGES, 1, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + weights = [(torch.randn(D, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + rotary = torch.randn(MAX_POSITIONS, 2, ROTARY // 2, device="cuda", generator=g) + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + return ( + packed, + kv, + None, + index_k, + slots, + None, + HEADS, + KV, + INDEX, + D, + ROTARY, + EPS, + *weights, + rotary, + positions, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _args(cell.spec["m"], zlib.crc32(cell.why.encode())) + certifying(op) + ref = op.reference(*args) + q, index_q = op(*args) + torch.cuda.synchronize() + _within_fp8(q.view(ref["q"].shape), ref["q"], "q") + _within_fp8(index_q.view(ref["index_q"].shape), ref["index_q"], "index_q") + kv, index_k, slots = args[1], args[3], args[4] + for t in range(cell.spec["m"]): + page, within = divmod(int(slots[t]), PAGE) + _within_fp8(kv[page, 0, :, within], ref["k"][t], f"k token {t}") + _within_fp8(kv[page, 1, :, within], ref["v"][t], f"v token {t}") + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..bf3c3c271148 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_blocks entry of the MiniMax-M3 tp8 target: MSA top-16 block selection with the forced +init / local blocks and per-row valid-block masks.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_blocks import ( # noqa: E501 + m3_select_blocks as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" + +TOPK, INIT, LOCAL = 16, 0, 1 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + n_blocks, total_q = spec["n_blocks"], spec["total_q"] + # Distinct scores (a permutation) so the selection has no ties to break. + scores = ( + torch.randperm(n_blocks * total_q, generator=g, device="cuda") + .float() + .view(1, n_blocks, total_q) + ) + if spec["mixed"]: + n_valid = torch.randint( + 17, n_blocks + 1, (total_q,), generator=g, device="cuda", dtype=torch.int32 + ) + else: + n_valid = torch.full((total_q,), n_blocks, dtype=torch.int32, device="cuda") + with validating(op): + out = op(scores, TOPK, n_valid, INIT, LOCAL, True) + ref = op.reference(scores, TOPK, n_valid, INIT, LOCAL, True) + assert out.shape == ref.shape + assert torch.equal(out, ref) + assert out.permute(1, 0, 2).is_contiguous(), "head-major backing" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..945025a650ca --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp8 (Rubin) m3_select_long entry: top-16 MSA block selection for long decode rows, 1 KV head.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long import ( # noqa: E501 + m3_select_long as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key +# widths, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 2048, 2049, 4096, 4097, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # re-read path). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + # Every row: 10 winners above, ~1500 blocks tied at the 16th score, the rest below; the in-tree op stages up to + # 2048 boundary candidates, so it must break this tie to the smaller block like this op. + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + # Odd tokens: almost everything -inf, so picks run past the finite scores. + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +def _mass_tie_rows(scores, nvb, init_blocks, local_blocks): + """[total_q, num_kv_heads] mask of rows where more than 2048 valid blocks share the 16th-largest score.""" + num_kv_heads, num_blocks, total_q = scores.shape + blocks = torch.arange(num_blocks, device=scores.device) + raw = nvb[:total_q].long() + valid = raw.clamp(0, num_blocks) + forced = scores.permute(2, 0, 1).clone() # [total_q, heads, blocks] + forced = torch.where(blocks < init_blocks, 1.0e30, forced) + forced = torch.where(blocks >= (raw - local_blocks).clamp_min(0)[:, None, None], 1.0e29, forced) + in_row = (blocks[None, :] < valid[:, None])[:, None, :] + ranked = torch.where(in_row, forced, float("-inf")).sort(dim=-1, descending=True).values + kth = ranked[..., min(TOPK, num_blocks) - 1 : min(TOPK, num_blocks)] + return ((forced == kth) & in_row).sum(dim=-1) > 2048 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + + # Selection is an exact claim: the same blocks or a different attention. + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # And it is the in-tree selector's answer, bit for bit (the op the target falls back to) -- except on rows + # where more than 2048 blocks tie at the 16th-largest score: there the in-tree op does not break the tie to + # the smallest blocks, while this op and the reference do. + scores, nvb, _, init_blocks, local_blocks, _ = args + in_tree = select_blocks_from_maxscore( + scores, + topk=TOPK, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + mass_tie = _mass_tie_rows(scores, nvb, init_blocks, local_blocks) + assert torch.equal(out[~mass_tie], in_tree[~mass_tie]), ( + "selected blocks differ from the in-tree minimax_m3_select_blocks" + ) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + """Releasing the dependent grid right after the select's own wait changes only when the next kernel may launch.""" + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + """Releasing the dependent grid once the row's scores are loaded changes only when the next kernel may launch.""" + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) + + +def test_guard_refuses_out_aliasing_scores() -> None: + scores, nvb, _, init_blocks, local_blocks, _ = _case(op.CELLS[0].spec, seed=3) + aliased = ( + scores.view(torch.int32) + .reshape(-1)[: scores.shape[2] * TOPK] + .view(scores.shape[2], 1, TOPK) + ) + with pytest.raises(AssertionError, match="alias"): + with validating(op): + op(scores, nvb, aliased, init_blocks, local_blocks, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py new file mode 100644 index 000000000000..d5af7b1a4668 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp8 (Rubin) m3_select_long_dsl entry: the long-row top-16 select from one 384-thread CTA per row, +against the reference and bit for bit against the 1024-thread m3_select_long on every cell.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( + m3_select_long_dsl as select_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, + m3_select_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long_dsl import ( # noqa: E501 + m3_select_long_dsl as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key +# widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 1536, 1537, 4224, 4225, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # whole-row re-rank). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # Every row, mass ties included, is the 1024-thread kernel's answer. + scores, nvb, _, init_blocks, local_blocks, _ = args + long_out = torch.empty_like(out) + m3_select_long(scores, nvb, long_out, init_blocks, local_blocks, False) + assert torch.equal(out, long_out), "selected blocks differ from m3_select_long" + + +@pytest.mark.parametrize("layout", ["kv_blocks_q", "kv_q_blocks"]) +@pytest.mark.parametrize("valid", [17, 1294, 2524, 3907, 9216]) +def test_layouts_and_widths_match_the_1024_thread_select(layout, valid) -> None: + """Both score layouts the scorer can hand over, at a valid count in each key width.""" + g = torch.Generator(device="cuda").manual_seed(valid) + nb = max(valid, 129) + s = torch.randn(1, nb, 8, generator=g, device="cuda") + if layout == "kv_q_blocks": + s = s.permute(0, 2, 1).contiguous().permute(0, 2, 1) + nvb = torch.tensor([valid - (t % 2) for t in range(8)], dtype=torch.int32, device="cuda") + out = torch.empty(1, 8, TOPK, dtype=torch.int32, device="cuda").permute(1, 0, 2) + ref = torch.empty_like(out) + op(s, nvb, out, 0, 1, False) + m3_select_long(s, nvb, ref, 0, 1, False) + assert torch.equal(out, ref) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +@pytest.mark.parametrize("early_trigger", [False, True]) +@pytest.mark.parametrize("cell", range(len(op.CELLS))) +def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: + """kernels/m3_select_long_dsl.SELECT_N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the + grid wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and + tied rows).""" + phase = 0 if early_trigger else 1 + spec = op.CELLS[cell].spec + scores, nvb, before, init_blocks, local_blocks, _ = _case(spec, seed=11 + cell) + after = torch.empty_like(before) + select_kernels._compile((True, early_trigger, phase, True))( + scores, nvb, before, init_blocks, local_blocks + ) + select_kernels._compile((True, early_trigger, phase, False))( + scores, nvb, after, init_blocks, local_blocks + ) + assert torch.equal(before, after) + + +def test_graph_replay_selects_the_same() -> None: + spec = op.CELLS[1].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=9) + ref = op.reference(scores, nvb, out, init_blocks, local_blocks) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + for _ in range(3): + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..3883cd543732 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,626 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode entry (the target's Triton block-sparse decode and its routes) at the TP8 heads. + +Besides the band of every route around a float64 reference, the Triton merge's MXFP8 o_proj hand-off (published to +``m3_prequant.SLOT``) is held to ``torch.ops.trtllm.mxfp8_quantize(output, True)`` bit for bit: on identity steps and on +selected steps with the long-context range killed (the routes the o_proj consumes when LC8 does not run). +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( + modeling, # noqa: E501 +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_attn_lc8, + m3_prequant, + m3_report, + triton_sparse_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode as stock_decode, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import ( # noqa: E402 — needs the path declared above + HEAD_DIM, + NUM_HEADS, + build_step, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +# Band, in BF16 ulps (2^-8) of each [8, 128] row's largest magnitude, and of the relative RMS. Every route +# computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference +# is the BF16 rounding on both sides (one ulp of the row at most), the FP16 probabilities before PV, the split-K +# partials kept in FP32 and merged by their log2-sum-exp, and approximate exp2/log2: 2 ulps per element, the two +# roundings (each at most half an ulp) dominating the RMS, so half an ulp. (GB300 TP4: at most 1.80 ulps per element +# and 0.20 ulps RMS over the cells.) +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, output, identity): + return op( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + identity, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=spec["identity"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + certifying(op) + _call(case, output, spec["identity"]) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + spec["identity"], + ) + assert torch.isfinite(output.float()).all(), "padding rows must be zeros, not NaN" + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_flag_does_not_change_the_answer() -> None: + """On an identity table the identity route (the per-row producer) and the selected-table route (the + long-context kernel) compute the same attention.""" + case = build_step((2048, 130), 4, identity=True, seed=17) + outs = [] + for identity in (True, False): + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, identity) + outs.append(output) + assert_within_ulp(outs[0], outs[1], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_guard_refuses_a_padded_selection_slot() -> None: + """The Triton producers index the block table with every one of a token's first min(16, blocks) ids, + -1 included.""" + case = build_step((1500,), 1, identity=True, seed=23) + case.topk[0, 0, 3] = -1 + output = torch.empty(1, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="blocks of the request"): + with validating(op): + _call(case, output, True) + + +def test_guard_refuses_a_false_identity_promise() -> None: + """identity_blocks lets the multi-query producer read one table for the whole request.""" + case = build_step((1500, 1600, 1700, 1800, 1500, 1600, 1700, 1800), 4, identity=True, seed=29) + case.topk[0, 5, :2] = case.topk[0, 5, :2].flip(0) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="identity_blocks promises"): + with validating(op): + _call(case, output, True) + + +def _swizzled_scale_offsets(rows: int, width: int) -> torch.Tensor: + """Offsets of rows [0, rows)'s UE8M0 scales in the 128x4-swizzled layout of mxfp8_quantize(x, True).""" + k_sf = width // 32 + k_tiles = (k_sf + 3) // 4 + r = torch.arange(rows, device="cuda").view(-1, 1) + k = torch.arange(k_sf, device="cuda").view(1, -1) + offsets = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + return offsets.reshape(-1) + + +def _mxfp8_handoff_matches_stock(output) -> None: + """The real rows' E4M3 bytes and UE8M0 scales equal mxfp8_quantize's. The merge writes only the step's rows into + its persistent 128-row scale tile; the padding rows' scales are never read by the o_proj and are not compared.""" + rows = output.shape[0] + flat = output.view(rows, -1) + taken = m3_prequant.SLOT.take(flat) + assert taken is not None, "the route published no MXFP8 hand-off for its output" + q8, sf = taken + stock_q, stock_sf = torch.ops.trtllm.mxfp8_quantize(flat, True) + assert torch.equal( + q8.view(torch.uint8).reshape(rows, -1), stock_q.view(torch.uint8).reshape(rows, -1) + ) + at = _swizzled_scale_offsets(rows, flat.shape[1]) + assert torch.equal( + sf.view(torch.uint8).reshape(-1)[at], stock_sf.view(torch.uint8).reshape(-1)[at] + ) + + +@pytest.mark.parametrize("seq_lens", [(1900,), (2048, 0), (700, 1500)]) +def test_identity_merge_mxfp8_handoff_is_stock_mxfp8_quantize(seq_lens) -> None: + """The Triton merge writes the o_proj hand-off (E4M3 rows, 128x4-swizzled UE8M0 scales) as mxfp8_quantize.""" + case = build_step(seq_lens, 4, identity=True, seed=31 + len(seq_lens)) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(case, output, True) + _mxfp8_handoff_matches_stock(output) + + +@pytest.mark.parametrize("seq_lens", [(40000,), (9000, 40000)]) +def test_killed_long_context_route(seq_lens) -> None: + """With SPARSE_DECODE_LC8 killed a selected table runs the Triton producer + merge: in the band of the reference, + and its MXFP8 hand-off is mxfp8_quantize of the rows it wrote.""" + case = build_step(seq_lens, 4, identity=False, seed=41 + len(seq_lens)) + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, False) + _mxfp8_handoff_matches_stock(output) + finally: + triton_sparse_decode.LC8_KEYS = landed + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_route_is_the_certified_producer_range() -> None: + """Every identity step at the TP8 heads (one KV head) takes the per-row producer below 17 rows and the multi-query + producer from 17 rows; modeling.SPARSE_DECODE_PRODUCER covers both, and the cells certify both.""" + assert modeling.SPARSE_DECODE_PRODUCER.covers("per_row") + assert modeling.SPARSE_DECODE_PRODUCER.covers("multi_query") + seen = set() + for cell in op.CELLS: + keys = tuple(modeling._sparse_decode_producer(cell.spec)) + rows = cell.spec["decode_query_len"] * len(cell.spec["seq_lens"]) + want = (("multi_query" if rows >= 17 else "per_row"),) if cell.spec["identity"] else () + assert keys == want, (cell.why, keys) + seen.update(keys) + assert seen == {"per_row", "multi_query"} + + +_IDENTITY_LENS = (1900, 2048, 700, 1500, 1200, 2000, 300, 1024) +_SELECTED_LENS = (9000, 40000, 20000, 5000, 7000, 12000, 30000, 3000) + + +def _stock(case, output): + stock_decode( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + sm_scale=case.sm_scale, + output=output, + decode_query_len=case.decode_query_len, + k_block_scale=case.k_block_scale, + v_block_scale=case.v_block_scale, + k_global_scale=case.k_global_scale, + v_global_scale=case.v_global_scale, + ) + + +@pytest.mark.parametrize("identity", [True, False], ids=["identity", "selected_lc_killed"]) +@pytest.mark.parametrize("requests", [1, 2, 4, 8]) +def test_triton_route_is_bitwise_the_stock_decode(requests, identity) -> None: + """The target's per-row Triton producer + split-K merge writes the stock in-tree MSA decode's rows bit for bit: on + SM107 it takes stock's split-K factor at every row count (4 / 8 / 16 / 32 verify rows here). An identity step of + at least SPARSE_DECODE_MQ_MIN_TOKENS rows takes the multi-query producer (one KV head), which stock does not run: + exact-class instead, in the band of the float64 reference and the same bits on every run.""" + lens = (_IDENTITY_LENS if identity else _SELECTED_LENS)[:requests] + case = build_step(lens, 4, identity=identity, seed=50 + requests + 100 * identity) + rows = case.q.shape[0] + multi_query = identity and rows >= modeling.SPARSE_DECODE_MQ_MIN_TOKENS + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + again = torch.empty_like(target) + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + _call(case, target, identity) + _call(case, again, identity) + finally: + triton_sparse_decode.LC8_KEYS = landed + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), again.view(torch.int16)), "run to run" + if not multi_query: + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)) + return + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + target, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + True, + ) + assert_within_ulp(target, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_plain_route_is_bitwise_the_stock_decode() -> None: + """Above MergeMxfp8.MAX_ROWS (here 160 verify rows) the Triton producer + merge writes no hand-off; with + SPARSE_DECODE_LC8 killed (its "plain" key covers these rows) it writes the stock decode's rows bit for bit.""" + case = build_step((3000,) * 40, 4, identity=False, seed=77) + target = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + _call(case, target, False) + finally: + triton_sparse_decode.LC8_KEYS = landed + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)) + + +@pytest.mark.parametrize("identity", [True, False], ids=["triton_identity", "lc8_selected"]) +def test_handoff_kill_publishes_nothing_and_keeps_the_rows(identity) -> None: + """MERGE_MXFP8_HANDOFF False: no route publishes the MXFP8 hand-off (the o_proj then quantizes its BF16 input), and + the attention rows are the same bits as with the hand-off.""" + case = build_step((1900,) if identity else (40000,), 4, identity=identity, seed=91 + identity) + on = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + off = torch.empty_like(on) + _call(case, on, identity) + assert m3_prequant.SLOT.take(on.view(on.shape[0], -1)) is not None + triton_sparse_decode.MERGE_MXFP8_HANDOFF = False + try: + _call(case, off, identity) + assert m3_prequant.SLOT.take(off.view(off.shape[0], -1)) is None + finally: + triton_sparse_decode.MERGE_MXFP8_HANDOFF = True + assert torch.equal(on.view(torch.int16), off.view(torch.int16)) + + +def _short_lens(requests: int, g: torch.Generator) -> tuple: + """Short contexts of 1..16 blocks of 128 tokens (the GSM / identity-route regime), a few exactly on a block edge.""" + lens = torch.randint(4, 16 * 128 + 1, (requests,), generator=g).tolist() + if requests > 1: + lens[0] = 128 * int(torch.randint(1, 17, (1,), generator=g)) + return tuple(lens) + + +@pytest.mark.parametrize("seed", range(12)) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_short_table_stress_is_bitwise_the_stock_decode(requests, seed) -> None: + """Short contexts (1..16 blocks) at 1..4 requests x 4 verify rows, the identity route and a random selected route + with -1 tails past each token's blocks: the target's Triton route writes the stock decode's rows bit for bit.""" + g = torch.Generator().manual_seed(1000 * requests + seed) + for identity in (True, False): + case = build_step(_short_lens(requests, g), 4, identity=identity, seed=7 * seed + requests) + target = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + _call(case, target, identity) + finally: + triton_sparse_decode.LC8_KEYS = landed + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + identity, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("fault", ["neg_in_window", "id_past_row", "page_out_of_pool"]) +def test_invalid_slot_fails_closed(fault) -> None: + """A corrupt decode input (a -1 inside a token's block window, an id past the block-table row, or a block-table + page outside the pool) is read as an absent block instead of a wild KV read: the rows equal the reference over the + remaining blocks.""" + case = build_step((1500, 2048, 700), 4, identity=False, seed=303) + tok, slot = 5, 2 + if fault == "neg_in_window": + case.topk[0, tok, slot] = -1 + elif fault == "id_past_row": + case.topk[0, tok, slot] = case.block_table.shape[1] + 7 + else: + req = tok // 4 + blk = int(case.topk[0, tok, slot]) + bad_page = case.k_paged.shape[0] + 1000 + case.block_table = case.block_table.clone() + case.block_table[req, blk] = bad_page + # Every token of the request that selected that block loses it. + hits = (case.topk[:, req * 4 : req * 4 + 4] == blk).nonzero().tolist() + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + _call(case, output, False) + finally: + triton_sparse_decode.LC8_KEYS = landed + torch.cuda.synchronize() + # The reference walks each token's first min(16, blocks) entries and drops -1: the corrupt slots as -1. + ref_topk = case.topk.clone() + if fault == "page_out_of_pool": + for h, t, k in hits: + ref_topk[h, req * 4 + t, k] = -1 + else: + ref_topk[0, tok, slot] = -1 + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + ref_topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert torch.isfinite(output.float()).all() + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def _identity_table( + seq_lens: torch.Tensor, decode_query_len: int, heads: int, lag: int = 0 +) -> torch.Tensor: + """The identity selection of every verify token, staged from lengths `lag` tokens behind seq_lens.""" + seq = (seq_lens.cpu().to(torch.int64) - lag).clamp(min=0).repeat_interleave(decode_query_len) + intra = torch.arange(seq.shape[0]) % decode_query_len + n_valid = ((seq - decode_query_len + intra + 1).clamp(min=0) + 127) // 128 + ar = torch.arange(16) + table = torch.where(ar[None, :] < n_valid[:, None], ar[None, :], -1).to(torch.int32) + return table.unsqueeze(0).expand(heads, -1, -1).contiguous().cuda() + + +# Each crosses a 128-token edge within its four verify tokens or between the stale and the true lengths. +_EDGE_LENS = (1284, 1287, 1281, 129, 132, 2048, 1028, 5, 643, 1920, 256, 1) + + +@pytest.mark.parametrize("fault", ["stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_identity_route_reads_its_own_lengths_not_the_table(requests, fault) -> None: + """The identity route derives block i for slot i from the kv lengths the kernel reads: a staged table that lags + the lengths by a verify step (a 128-token edge crossed), or is garbage, gives the stock decode's rows over the + true identity selection bit for bit.""" + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=400 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + assert torch.equal(case.topk, true_table), ( + "build_step's identity table is the identity selection" + ) + if fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _stock(case, stock) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + # The Triton identity route (triton_sparse_decode.LC8_IDENTITY off); the LC8 identity route has its own test below. + triton_sparse_decode.LC8_IDENTITY = False + try: + _call(case_staged, target, True) + finally: + triton_sparse_decode.LC8_IDENTITY = True + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("fault", ["none", "stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_lc8_identity_route_reads_its_own_lengths_not_the_table( + requests, fault, monkeypatch +) -> None: + """triton_sparse_decode.LC8_IDENTITY: an identity step runs the LC8, which takes block i for CTA i from the kv + lengths it reads. Whatever the staged table holds (true, lagging a verify step, all -1, garbage), the rows and the + MXFP8 hand-off are the LC8's on the true identity table (the selected-table route) bit for bit, and in the band of + the float64 reference.""" + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=700 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + if fault == "none": + staged = true_table.clone() + elif fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + selected = torch.empty_like(target) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + case_true = SimpleNamespace(**{**vars(case), "topk": true_table}) + taken = [] + monkeypatch.setattr(m3_report, "taken", lambda path, rows, what="": taken.append(path)) + _call(case_staged, target, True) + monkeypatch.undo() + # The LC8 engagement, plus its one-division merge's (m3_attn_lc8.ONE_DIV_MERGE) when on. + assert sorted(taken) == sorted( + [triton_sparse_decode.LC8_PATH] + + ([triton_sparse_decode.ONE_DIV_MERGE_PATH] if m3_attn_lc8.ONE_DIV_MERGE else []) + ), taken + target_handoff = m3_prequant.SLOT.take(target.view(rows, -1)) + _call(case_true, selected, False) + selected_handoff = m3_prequant.SLOT.take(selected.view(rows, -1)) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), selected.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) + assert (target_handoff is None) == (selected_handoff is None) + if target_handoff is not None: + for a, b in zip(target_handoff, selected_handoff): + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + true_table, + case.block_table, + case.seq_lens, + case.sm_scale, + target, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert_within_ulp(target, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +@pytest.mark.parametrize("requests", [1, 2]) +def test_lc8_identity_step_reads_q_after_the_wait(requests, monkeypatch) -> None: + """An identity step has no select, so the LC8 must read q after its grid wait (783877: a route that skipped it + read q before the producer wrote it): the op gets identity and identity_step on every identity step, and the rows + are the same call after call.""" + lens = tuple(_EDGE_LENS[(requests * 3 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=900 + requests) + true_table = _identity_table(case.seq_lens, 4, case.topk.shape[0]) + case_true = SimpleNamespace(**{**vars(case), "topk": true_table}) + rows = case.q.shape[0] + seen = [] + real = m3_attn_lc8.m3_sparse_decode_lc8 + + def spy(*args, **kwargs): + seen.append((kwargs.get("identity", False), kwargs.get("identity_step", False))) + return real(*args, **kwargs) + + monkeypatch.setattr(m3_attn_lc8, "m3_sparse_decode_lc8", spy) + out = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + ref = torch.empty_like(out) + _call(case_true, out, True) + _call(case_true, ref, True) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert seen[0] == (True, True), seen + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)) + + +@pytest.mark.parametrize("requests", [2, 3]) +def test_lc_fails_closed_on_ids_past_the_row(requests) -> None: + """A selected id past the token's row (a stale table, e.g. a padded row keeping a 500K request's n_valid) is + skipped, as the Triton producer skips it: the rows equal those with that slot at -1, bit for bit, and stay finite + even though the block-table tail it would reach points at a page whose block scales are NaN (0xFF).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_prequant, + ) + + lens = (165521, 700, 3000)[:requests] + case = build_step(lens, 4, identity=False, seed=1300 + requests) + dql = case.decode_query_len + good = case.topk.clone() + bad = case.topk.clone() + for tok in range(case.q.shape[0]): + req, qoff = divmod(tok, dql) + kv_len = max(int(case.seq_lens[req]) - dql + qoff + 1, 0) + nblk = (kv_len + 127) // 128 + past = nblk + 1 + if 2 <= nblk < 16 and past < case.block_table.shape[1]: + bad[:, tok, nblk - 1] = past + good[:, tok, nblk - 1] = -1 + page = int(case.block_table[req, past]) + case.k_block_scale[page].fill_(0xFF) + case.v_block_scale[page].fill_(0xFF) + assert not torch.equal(bad, good) + outs = [] + for table in (bad, good): + out = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(SimpleNamespace(**{**vars(case), "topk": table}), out, False) + outs.append(out) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert torch.isfinite(outs[0].float()).all() + assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..cc8de6c62717 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode_lc8 entry (one-launch long-context sparse decode at the TP8 8 query heads, +``torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8``). + +Every cell is held to the float64 reference within the derived band, its MXFP8 hand-off to the recipe bit for bit, and +the two ``prewait`` variants (staging before or after the grid wait) to each other bit for bit. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_sparse_decode_lc8 import ( # noqa: E501 + NUM_HEADS, + WIDTH, + dequant_nvfp4_pages, + mxfp8_handoff_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_sparse_decode_lc8 import ( # noqa: E501 + m3_sparse_decode_lc8 as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import NUM_HEADS as CASE_HEADS # noqa: E402 — needs the path declared above +from _m3_nvfp4_cache import build_step # noqa: E402 +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_sparse_decode_lc8 requires a CUDA device" + +#: Hand-off buffers are sized as MERGE_MXFP8 sizes them. +_HANDOFF_ROWS = 128 + +# Band, in BF16 ulps (2^-8) of each output row's largest magnitude, and of the relative RMS. +# +# The reference is float64 attention over the exactly dequantized blocks, rounded once to BF16. The kernel's +# operands are exact too (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so what separates the two is: the final +# BF16 rounding on both sides (up to one ulp of the element, at most one of the row's scale), the softmax +# probabilities rounded to FP16 before PV (2^-11 relative each, averaging down over the block), ex2.approx / +# lg2.approx (~2^-22) and FP32 accumulation order. Hence 2 ulps per element; the RMS is dominated by the two +# roundings (each at most half an ulp, about 0.3 in aggregate), so half an ulp. None of these depends on the head +# count or on the MMA orientation: the tp4 target's 16-head kernel measured at most 1.91 ulps per element and 0.25 +# ulps RMS on GB300 over the same cells. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _args(case, out, q8, sf, linear_softmax, prewait, num_sf_k_tiles=WIDTH // 128): + return ( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + case.decode_query_len, + case.sm_scale, + linear_softmax, + False, + prewait, + ) + + +def _buffers(total_q, mxfp8): + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + if not mxfp8: + return out, None, None + q8 = torch.zeros(_HANDOFF_ROWS, WIDTH, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_HANDOFF_ROWS * (WIDTH // 32), dtype=torch.uint8, device="cuda") + return out, q8, sf + + +def _case(spec, seed): + return build_step(spec["seq_lens"], spec["decode_query_len"], identity=False, seed=seed) + + +def test_the_cases_are_the_tp8_geometry() -> None: + assert CASE_HEADS == NUM_HEADS == 8 and WIDTH == 1024 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = _case(spec, zlib.crc32(cell.why.encode()) % 2**31) + total_q = case.q.shape[0] + out, q8, sf = _buffers(total_q, spec["mxfp8"]) + args = _args(case, out, q8, sf, spec["linear_softmax"], spec["prewait"]) + certifying(op) + op(*args) + ref = op.reference(*args) + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + # Padding rows attend nothing and must be exact zeros, not NaN. + assert torch.isfinite(out.float()).all() + if spec["mxfp8"]: + # The hand-off is a deterministic function of the op's own BF16 rows: exact. + ref_q8, ref_sf = mxfp8_handoff_reference(out, _HANDOFF_ROWS) + assert torch.equal(q8[:total_q], ref_q8), ( + "MXFP8 rows differ from the recipe applied to the BF16 rows" + ) + assert torch.equal(sf, ref_sf), "UE8M0 scales differ from the recipe / swizzle" + + # The other staging variant: the same arithmetic, so the same bits (rows and hand-off). + other_out, other_q8, other_sf = _buffers(total_q, spec["mxfp8"]) + op(*_args(case, other_out, other_q8, other_sf, spec["linear_softmax"], not spec["prewait"])) + assert torch.equal(other_out.view(torch.int16), out.view(torch.int16)), ( + "prewait and plain staging wrote different rows" + ) + if spec["mxfp8"]: + assert torch.equal(other_q8, q8) and torch.equal(other_sf, sf) + + +def test_reference_dequant_matches_the_quantizer() -> None: + """The reference reads the paged bytes through its own layout claims (K scales token-major, V scales in + 4x4 token quads). Hold them to the quantizer's own flat output: every slot was written in order, so slot + ``s`` is page ``s // 128``, token ``s % 128``, and the flat scales are token-major by construction.""" + case = build_step((3000,), 1, identity=False, seed=5) + pages = torch.arange(case.k_paged.shape[0], device="cuda") + views = ((case.k_paged, case.k_block_scale, False), (case.v_paged, case.v_block_scale, True)) + for role, (data, scale, quad) in enumerate(views): + flat_data, flat_scale = torch.ops.trtllm.fp4_quantize( + case.source[role].reshape(case.num_slots, 1, -1).contiguous(), + case.quant[role + 1 : role + 2], + 16, + False, + False, + ) + expected = dequant_nvfp4_pages( + flat_data.view(torch.uint8).reshape(-1, 1, 128, 64), + flat_scale.view(torch.uint8).reshape(-1, 1, 128, 8), + pages, + quad_scales=False, + ) + got = dequant_nvfp4_pages(data, scale, pages, quad_scales=quad) + assert torch.equal(got, expected), ( + f"role {role}: paged dequantization disagrees with the quantizer" + ) + + +def test_the_variants_agree() -> None: + """The two softmax recurrences, the hand-off switch and the staging must not change the BF16 rows beyond + rounding; the hand-off switch and the staging not at all.""" + case = build_step((5000, 12345), 4, identity=False, seed=11) + outs = {} + for linear in (False, True): + for mxfp8 in (False, True): + for prewait in (False, True): + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op(*_args(case, out, q8, sf, linear, prewait)) + outs[(linear, mxfp8, prewait)] = out + base = outs[(False, False, False)] + for key, other in outs.items(): + assert_within_ulp(other, base, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + same_recurrence = outs[(key[0], False, False)] + torch.testing.assert_close(other, same_recurrence, rtol=0, atol=0) + + +def test_guard_refuses_a_block_past_the_request() -> None: + """A selected id at or past the request's block count reads a page the token cannot attend; with every + position masked the CTA's softmax is NaN.""" + case = build_step((3000,), 1, identity=False, seed=4) + case.topk[0, 0, 0] = 40 + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="block count"): + with validating(op): + op(*_args(case, out, q8, sf, False, True)) + + +def test_guard_refuses_strided_page_rows() -> None: + """The kernel walks a page's token rows at a fixed 64-byte pitch; a padded row stride is misread.""" + case = build_step((3000,), 1, identity=False, seed=6) + wide = torch.zeros(*case.k_paged.shape[:3], 80, dtype=torch.uint8, device="cuda")[..., :64] + wide.copy_(case.k_paged) + case.k_paged = wide + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="64-byte rows"): + with validating(op): + op(*_args(case, out, q8, sf, False, True)) + + +def test_guard_refuses_a_narrow_handoff() -> None: + """The hand-off rows are 1024 wide (8 scale k-tiles); a smaller count overlaps them.""" + case = build_step((3000,), 1, identity=False, seed=7) + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="num_sf_k_tiles"): + with validating(op): + op(*_args(case, out, q8, sf, False, True, num_sf_k_tiles=4)) + + +@pytest.mark.parametrize("prewait", [True, False], ids=["prewait", "plain"]) +@pytest.mark.parametrize("mxfp8", [True, False], ids=["handoff", "no_handoff"]) +def test_run_to_run_deterministic(prewait, mxfp8) -> None: + """The exact-class reference for the TP8 sparse decode (band vs fp64, as the TP4 LC): every launch of the same + step writes the same bits, in eager repeats and in a CUDA-graph replay.""" + spec = dict(seq_lens=(165521, 5000), decode_query_len=4) + case = _case(spec, seed=91) + for linear in (False, True): + runs = [] + for _ in range(3): + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op(*_args(case, out, q8, sf, linear, prewait)) + runs.append((out, q8, sf)) + g_out, g_q8, g_sf = _buffers(case.q.shape[0], mxfp8) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*_args(case, g_out, g_q8, g_sf, linear, prewait)) + graph.replay() + torch.cuda.synchronize() + runs.append((g_out, g_q8, g_sf)) + base = runs[0] + for other in runs[1:]: + assert torch.equal(other[0].view(torch.int16), base[0].view(torch.int16)), ( + linear, + prewait, + mxfp8, + ) + if mxfp8: + assert torch.equal(other[1], base[1]) and torch.equal(other[2], base[2]), ( + linear, + prewait, + ) + + +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + [((165521,), 4), ((5000, 12345), 4), ((700,), 1), ((200000, 3000, 64, 70000), 4)], + ids=["1x165k", "2req", "short-1row", "4req-incl-64"], +) +def test_merge_once_is_bitwise_the_image_op(seq_lens, decode_query_len) -> None: + """kernels/m3_attn_lc8.MERGE_ONCE: the target's op (merge weights formed once per CTA) writes exactly the + image op's rows and hand-off, at both softmax recurrences and both staging variants, including requests with fewer + than 16 selected blocks (empty chunks, -inf LSE). With SPARSE_DECODE_LC8_PREFETCH_OPROJ it also pulls two ranges + the size of the TP8 o_proj weight and scales (one not a multiple of 16 bytes) into L2: the same bits again. With + SPARSE_DECODE_LC8_PREFETCH_NEXT_KV it records each row's selection (the first min(16, blocks) ids, -1 after) and + prefetches a next layer's recorded blocks, stale ids past a request's blocks included: the same bits, and the + record.""" + assert hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_sparse_decode_lc8"), ( + "the target's LC8 op is not loaded (its TP8 extension)" + ) + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + ) + o_proj_weight = torch.empty(6144, 1024, dtype=torch.uint8, device="cuda") + o_proj_scale = torch.empty(6144 * 32 + 9, dtype=torch.uint8, device="cuda") + rows = case.q.shape[0] + n1_next = torch.randint( + -1, 2 * (max(seq_lens) // 128 + 1), (128, 16), dtype=torch.int32, device="cuda" + ) + for linear in (False, True): + for prewait in (False, True): + for mxfp8 in (False, True): + n1_self = torch.full((128, 16), -7, dtype=torch.int32, device="cuda") + next_kv = ( + n1_self, + n1_next, + case.k_paged.view(torch.uint8), + case.v_paged.view(torch.uint8), + case.k_block_scale.view(torch.uint8), + case.v_block_scale.view(torch.uint8), + ) + variants = ( + (False, None, None), + (True, None, None), + (True, (o_proj_weight, o_proj_scale), None), + (True, (o_proj_weight, o_proj_scale), next_kv), + ) + got = [] + for merge_once, prefetch, n1 in variants: + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op( + *_args(case, out, q8, sf, linear, prewait), + merge_once=merge_once, + prefetch=prefetch, + next_kv=n1, + ) + got.append((out, q8, sf)) + torch.cuda.synchronize() + a, aq, asf = got[0] + for i, (b, bq, bsf) in enumerate(got[1:], 1): + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), ( + linear, + prewait, + mxfp8, + i, + ) + if mxfp8: + assert torch.equal(aq, bq) and torch.equal(asf, bsf), (linear, prewait, i) + # The record: each row's first min(16, blocks) selected ids, -1 past them; rows past the step untouched. + want = case.topk[0].clone() + for tok in range(rows): + req, qoff = divmod(tok, case.decode_query_len) + kv_len = max(int(case.seq_lens[req]) - case.decode_query_len + qoff + 1, 0) + want[tok, min(16, (kv_len + 127) // 128) :] = -1 + assert torch.equal(n1_self[:rows], want), "the recorded selection" + assert bool((n1_self[rows:] == -7).all()), "rows past the step" + + +@pytest.mark.parametrize("linear", [False, True]) +@pytest.mark.parametrize("prewait", [True, False], ids=["prewait", "plain"]) +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + [((165521,), 4), ((5000, 12345), 4), ((700,), 1), ((200000, 3000, 64, 70000), 4)], + ids=["1x165k", "2req", "short-1row", "4req-incl-64"], +) +def test_one_div_merge_is_deterministic_and_in_band( + seq_lens, decode_query_len, prewait, linear +) -> None: + """kernels/m3_attn_lc8.ONE_DIV_MERGE (numerics rule 1b, gate a): the one-division merge writes the same rows on + every run (5 eager runs, PDL on and off) and under CUDA-graph replay (3 replays == eager), MXFP8 hand-off included; + the rows are within the float64 reference's band, and within summation-order distance of the per-chunk merge.""" + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + 3 + ) + total_q = case.q.shape[0] + + def run(out, q8, sf, use_pdl, one_div): + args = list(_args(case, out, q8, sf, linear, prewait)) + args[18] = use_pdl + op(*args, merge_once=True, one_div=one_div) + + runs = [] + for i in range(6): + out, q8, sf = _buffers(total_q, True) + run(out, q8, sf, i == 5, True) + runs.append((out, q8, sf)) + torch.cuda.synchronize() + a, aq, asf = runs[0] + for b, bq, bsf in runs[1:]: + assert torch.equal(a.view(torch.int16), b.view(torch.int16)) + assert torch.equal(aq, bq) and torch.equal(asf, bsf) + out = torch.empty_like(a) + run(out, None, None, True, True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(out, None, None, True, True) + for _ in range(3): + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), a.view(torch.int16)) + per, _, _ = _buffers(total_q, False) + run(per, None, None, False, False) + ref = op.reference(*_args(case, a, None, None, linear, prewait)) + assert torch.isfinite(a.float()).all() + assert_within_ulp(a, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + assert_within_ulp(a, per, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py new file mode 100644 index 000000000000..769a903d51c2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for kernels/m3_tokmap_guard: a one-request multi-token call of +thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a +stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the +appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for +every length. Without the guard the stale map moves rows (the negative control), and a several-request call passes +through with its own map untouched. + +The fresh-map reference forces the op's own build (a scaled RoPE type with no inv-freq cache, at rotary dim 0, so no +row is rotated either way). +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_tokmap_guard as guard, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = 8, 1, 128, 128 +BUCKET = 4 +_PEER_EXTRA = 777 # the second request's extra length behind the first +FP8_KV_QUANT_MODE = 128 +_DEV = torch.device("cuda") # unindexed on purpose: install() and the lookup normalize it +_TAKEN = [] + + +@pytest.fixture(scope="module", autouse=True) +def _installed(): + guard.install(_DEV, 48, _TAKEN.append) + yield + + +class _Step: + """One request (or ``nreq``) of BUCKET padded rows over a fresh FP8 paged cache, persistent across calls like a + real step's buffers.""" + + def __init__(self, kv_len: int, nreq: int = 1, seed: int = 0): + g = torch.Generator(device=_DEV).manual_seed(seed + kv_len) + self.nreq = nreq + # Per-request inputs (page tables, generation lengths, position offsets) for at least two requests, the page + # tables at the peer's longer length (_seq_with_peer): the stale map sends rows to request 1, which must have + # real entries behind it, also on the unguarded calls. + tables = max(nreq, 2) + pages = math.ceil((kv_len + _PEER_EXTRA + BUCKET) / PAGE) + self.nblocks = tables * pages + 2 + self.pool0 = ( + torch.randn(self.nblocks, 2, KVH, PAGE, D, device=_DEV, generator=g) + .clamp_(-4, 4) + .to(torch.float8_e4m3fn) + ) + self.pool = self.pool0.clone() + perm = torch.randperm(self.nblocks, device=_DEV, generator=g)[: tables * pages].view( + tables, pages + ) + self.offsets = torch.zeros(1, tables, 2, pages, dtype=torch.int32, device=_DEV) + self.offsets[0, :, 0] = 2 * perm + self.offsets[0, :, 1] = 2 * perm + 1 + self.pointers = torch.tensor([[self.pool.data_ptr(), 0]], dtype=torch.int64) + self.mapping = torch.zeros(1, 2, dtype=torch.int32) + self.qkv = torch.randn( + nreq * BUCKET, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=_DEV, generator=g + ) + # Request lengths after this step's append; generation lengths (accepted rows) set per case. + self.seq = torch.full((nreq,), kv_len, dtype=torch.int32, device=_DEV) + self.gen = torch.full((tables,), BUCKET, dtype=torch.int32, device=_DEV) + self.pos = ( + torch.arange(BUCKET, dtype=torch.int32, device=_DEV).repeat(tables, 1).contiguous() + ) + self.one = torch.ones(1, dtype=torch.float32, device=_DEV) + self.ws = torch.empty(64 << 20, dtype=torch.uint8, device=_DEV) + self.kv_len = kv_len + + def call(self, op, fresh_map: bool = False): + """One generation preprocess; returns a copy of the processed query.""" + out = op( + self.qkv, + self.ws, + self.seq, + self.gen, + self.pos, + self.offsets, + self.pointers, + self.mapping, + self.one, + self.one, + None, + None, + None, + None, + 0, + 0, + H, + KVH, + D, + PAGE, + FP8_KV_QUANT_MODE, + 1 << 20, + 1 << 20, + self.nreq * BUCKET, + self.nreq, + BUCKET, + self.kv_len, + 0, + 10000.0, + 1 if fresh_map else 0, + 1.0, + 1 << 20, + 0, + 1.0, + 1.0, + False, + BUCKET, + 0, + torch.cuda.get_device_properties(0).multi_processor_count, + self.nblocks, + 2, + True, + False, + ) + return out[0].clone() + + def poison_map(self) -> None: + """A stale multi-request map in the op's one-request map slice: rows 1.. on request 1.""" + layout = thop.get_trtllm_gen_generation_workspace_layout( + self.qkv.dtype, 1, self.nreq * BUCKET, H, D, 0, KVH, 0, False, False + ) + off = int(layout["tokens_info_offset"]) + stale = torch.tensor( + [[0, 0]] + [[1, t] for t in range(BUCKET - 1)], dtype=torch.int32, device=_DEV + ) + self.ws[off : off + 8 * BUCKET].view(torch.int32).view(BUCKET, 2).copy_(stale) + + def reset(self) -> None: + self.pool.copy_(self.pool0) + + +def _seq_with_peer(step: _Step) -> None: + # A second request's length behind the first, so a row the stale map moves to request 1 reads a real length. + step.seq = torch.tensor( + [step.kv_len, step.kv_len + _PEER_EXTRA], dtype=torch.int32, device=_DEV + ) + + +@pytest.mark.parametrize("kv_len", [130, 8191, 34117]) +@pytest.mark.parametrize("n", [1, 2, 3, 4]) +def test_guard_equals_a_fresh_map(kv_len, n) -> None: + step = _Step(kv_len) + _seq_with_peer(step) + step.gen.fill_(n) + step.poison_map() + step.reset() + q_ref = step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + + step.poison_map() + step.reset() + q = step.call(thop.trtllm_gen_generation_preprocess) + # The processed query's n live rows (rows past the request's length are padding the op leaves unwritten). + assert torch.equal(q[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)) + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + assert BUCKET in _TAKEN + + +def test_stale_map_moves_rows_without_the_guard() -> None: + step = _Step(34117) + _seq_with_peer(step) + step.gen.fill_(BUCKET) + step.poison_map() + step.reset() + step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + step.poison_map() + step.reset() + step.call(guard._original) + assert not torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + + +def test_one_captured_graph_serves_every_length() -> None: + step = _Step(34000) + _seq_with_peer(step) + refs = {} + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + q = step.call(guard._original, fresh_map=True) + refs[n] = (q, step.pool.clone()) + step.gen.fill_(BUCKET) + step.call(thop.trtllm_gen_generation_preprocess) # eager first, as the warmup runs it + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = thop.trtllm_gen_generation_preprocess + q_static = step.call(out) + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + g.replay() + torch.cuda.synchronize() + q_ref, pool_ref = refs[n] + assert torch.equal(q_static[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)), n + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)), n + + +def test_several_requests_pass_through() -> None: + step = _Step(8191, nreq=2) + layout = thop.get_trtllm_gen_generation_workspace_layout( + step.qkv.dtype, 2, 2 * BUCKET, H, D, 0, KVH, 0, False, False + ) + step.reset() + step.call(thop.trtllm_gen_generation_preprocess) + pool_guarded = step.pool.clone() + step.reset() + step.call(guard._original) + assert torch.equal(step.pool.view(torch.uint8), pool_guarded.view(torch.uint8)) + assert int(layout["tokens_info_size"]) >= 8 * 2 * BUCKET diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..a9d64f3a87f7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the msa_kv_update entry (the one-launch MSA KV-length correction).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.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_107__tp8.ops.attention.msa_kv_update import ( # noqa: E501 + TOKENS_PER_BLOCK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" + +#: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. +_SPARE = 3 + + +def _case(spec, seed): + """The metadata's buffers for one step, with stale values in every output so a skipped write shows.""" + g = torch.Generator(device="cpu").manual_seed(seed) + batch, qlen = spec["batch"], spec["query_len"] + total_q = batch * qlen + staged = torch.randint(qlen, 5000, (batch,), generator=g, dtype=torch.int32) + accept = spec["accept"] + if accept == "full": + kv = staged.clone() + elif accept == "page_edge": + staged = torch.tensor([130, 257, 2 * TOKENS_PER_BLOCK + qlen][:batch], dtype=torch.int32) + kv = staged - torch.tensor([0, 1, qlen - 1][:batch], dtype=torch.int32) + else: + kv = staged - torch.randint(0, qlen, (batch,), generator=g, dtype=torch.int32) + if accept == "padded": + kv[batch // 2 :] = 0 + staged[batch // 2 :] = 0 + ncols = (5000 + qlen) // TOKENS_PER_BLOCK + 2 + table = ( + torch.randperm(batch * ncols + 7, generator=g)[: batch * ncols] + .to(torch.int32) + .reshape(batch, ncols) + ) + dev = dict(device="cuda", dtype=torch.int32) + stale = -12345 + return ( + torch.cat([kv, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.cat([staged, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.full((batch + _SPARE,), qlen, **dev), + torch.arange(batch, **dev).repeat_interleave(qlen), + torch.arange(qlen, **dev).repeat(batch), + table.cuda(), + TOKENS_PER_BLOCK, + batch, + total_q, + torch.full((batch + _SPARE,), stale, **dev), + torch.full((batch + 1 + _SPARE,), stale, **dev) if spec["cu_kv"] else None, + torch.full((total_q + _SPARE,), stale, **dev), + torch.full((total_q + _SPARE,), stale, **dev) if spec["n_valid"] else None, + ) + + +def _check(args, ref): + seq, cu_kv, loc, n_valid = args[9], args[10], args[11], args[12] + batch, total_q = args[7], args[8] + ref_seq, ref_cu, ref_loc, ref_nv = ref + # Integer bookkeeping: exact, and nothing past the live rows is touched. + assert torch.equal(seq[:batch], ref_seq) + assert torch.equal(loc[:total_q], ref_loc) + assert bool((seq[batch:] == -12345).all()) and bool((loc[total_q:] == -12345).all()) + assert (cu_kv is None) == (ref_cu is None) and (n_valid is None) == (ref_nv is None) + if cu_kv is not None: + assert torch.equal(cu_kv[: batch + 1], ref_cu) + assert bool((cu_kv[batch + 1 :] == -12345).all()) + if n_valid is not None: + assert torch.equal(n_valid[:total_q], ref_nv) + assert bool((n_valid[total_q:] == -12345).all()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + op(*args) + _check(args, op.reference(*args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_hook_takes_every_cell(cell) -> None: + """The target's hook (what the metadata calls) takes each certified step and writes what the op does.""" + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + 1) + assert hook(*args) is True + _check(args, op.reference(*args)) + + +def test_the_hook_declines_what_the_op_would_misread() -> None: + """A strided operand is read wrongly by the op, so the hook declines it (the metadata then runs its own + torch update) and writes nothing.""" + args = list(_case(op.CELLS[2].spec, seed=5)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + before = args[11].clone() + assert hook(*args) is False + assert torch.equal(args[11], before) + + +def test_guard_refuses_a_strided_operand() -> None: + args = list(_case(op.CELLS[2].spec, seed=6)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_index_past_the_batch() -> None: + args = list(_case(op.CELLS[2].spec, seed=7)) + args[3] = args[3].clone() + args[3][-1] = args[7] + with pytest.raises(AssertionError, match="q_batch_row"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py new file mode 100644 index 000000000000..45b043d24e95 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py @@ -0,0 +1,454 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the ``m3_allreduce`` entry, 2 ranks (SM107): its cells against the entry's reference, +and the fused draft loop's layer boundaries built on it (``M3Eagle3Draft.ar_norm``: the plain all-reduce, then +FlashInfer's fused add + RMSNorm) against stock Eagle3's, bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy a single-node job can run (AUTO, ONESHOT, +TWOSHOT, NCCL; MNNVL, which the tp8 drafter's two-tray group may take, is multi-node only and not driven here), at the +drafter's shapes (H = 6144, eps = 1e-6, 1 to 32 rows), at the job's world size (2, 4 or 8 ranks on one node; the tp8 +drafter runs at 8 over two trays, which one node cannot form, and nothing here depends on the rank count). The inputs +are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own launcher, one rank per device in CUDA_VISIBLE_DEVICES, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the +collected entry point is ``test_modeling_v2_m3_allreduce.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py +""" + +import math +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 4, 8, 32) +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +#: World sizes the matrix runs at (one node): see ``_rank_job.WORLD_SIZES``. +WORLD_SIZES = (2, 4, 8) +RANK = 0 +WORLD = 1 + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _spread(parts: List[torch.Tensor]) -> List[torch.Tensor]: + """Every rank's BF16 partial: ``parts`` on the first ranks, zeros on the rest.""" + assert len(parts) <= WORLD, (len(parts), WORLD) + zeros = [torch.zeros_like(parts[0]) for _ in range(WORLD - len(parts))] + return [p.bfloat16() for p in parts] + [z.bfloat16() for z in zeros] + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on rank 0, x/4 on ranks 1 and 2 (x/2 on rank 1 at two ranks), + # 0 on the rest. + parts = [x / 2, x / 2] if WORLD == 2 else [x / 2, x / 4, x / 4] + return ( + _spread(parts), + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A, -A, +B, -B + d; +A, -A + d at two or three ranks) onto a unit + residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a, -a, b, (-b + d)] if WORLD >= 4 else [a, (-a + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return _spread(parts), residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +def _ulps_at(out: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the magnitude ``scale`` (the operands'), not at the result's.""" + return ((out.double() - ref.double()).abs() / _ulp_at(scale.abs())).max().item() + + +def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary (``M3Eagle3Draft.ar_norm``): ``m3_allreduce`` plain, then the catalog's + ``flashinfer_fused_add_rmsnorm``.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + + return flashinfer_fused_add_rmsnorm(m3_allreduce(x, ar, None), residual.clone(), w, EPS) + + +def _check_cells(mapping: Any) -> None: + """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum + to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude + u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final + roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). + The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm.functional import AllReduceStrategy + + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating + + ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + for i, cell in enumerate(m3_allreduce.CELLS): + spec = cell.spec + xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) + params = ( + None + if spec.get("plain") + else AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual.clone(), + norm_weight=w, + eps=EPS, + trigger_completion_at_end=False, + ) + ) + # The reference first: the fused path may write the updated residual into ``params.residual``. + ref = m3_allreduce.reference(xs, ar, params) + with validating(m3_allreduce): + out = m3_allreduce(xs[RANK], ar, params) + torch.cuda.synchronize() + total = sum(x.double() for x in xs) + if params is None: + err = _ulps_at(out, ref, total) + assert err <= 1.0, (cell.why, err) + else: + err = _ulps_at(out[1], ref[1], torch.maximum(total.abs(), residual.double().abs())) + assert err <= 2.5, (cell.why, "residual", err) + h = out[1].double() + own = h * torch.rsqrt(h.square().mean(-1, keepdim=True) + EPS) * w.double() + assert _ulps(out[0], own) <= 2.0, (cell.why, "normed", _ulps(out[0], own)) + dist.barrier() + if RANK == 0: + print(f"{RESULT_TAG} cell {cell.why} passed", flush=True) + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + _check_cells(mapping) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + ulp, res_ulp = 0.0, 0.0 + for m in TOKENS: + xs, residual, w = build(m) + target_out, target_res = _target(ar, xs[RANK], residual, w) + stock_out, stock_res = _stock(ar, xs[RANK], residual, w) + torch.cuda.synchronize() + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {name} m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {name:34s} | target = stock vs fp64: normed {ulp:.3f}" + f" | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD in WORLD_SIZES, f"the matrix runs at {WORLD_SIZES} ranks on one node, got {WORLD}" + try: + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns (one rank each)" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size in WORLD_SIZES, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at {WORLD_SIZES}" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..abf7fe23ea60 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,788 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``, at the job's ranks. + +Its own launcher, as the catalog's collective matrices are: run plainly it spawns +one rank process per device in CUDA_VISIBLE_DEVICES and kills them all on the +first failure or on the deadline. A rank whose peers never push polls forever, +so the deadline is what keeps a broken exchange from wedging the calling run. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_mc_allreduce_norm_op_matrix.py + +World size: one rank per visible device, 2, 4 or 8. The tp8 target runs these +entries at 8 ranks over two trays, which one 4-device node cannot form; the +world size is a runtime argument of the kernel and no cell's spec names one, so +the cells are certified at the world size of the job (``_rank_job.WORLD_SIZES``). + +Op: the entries call the image's op by default and this target's own copy +(``trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm``) with ``own_op``, which is +what the target's boundaries call when it is loaded. Every cell is driven +through each op the build has. + +Each rank builds the three multicast workspaces (8 / 16 / 48 tokens) the target +builds, over the torch.distributed group (the non-MPI path of +``_make_mnnvl_mcast_buffer``). Every rank can regenerate every rank's inputs from +the shared seed convention, which is what keeps each reference arithmetic rather +than a second collective. Producer pushes (``prepushed`` 2 / 3) are emulated by +storing the canonicalized rows through the multicast mapping at the call parity +the flags select, which is exactly what the producing kernels store. + +Not a pytest module, despite the ``check_*`` bodies: they are one fixed sequence +inside a single multi-rank job whose workspace state (call parities) carries from +one to the next. The collected entry point is +``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +finalize_entry: Any = None +ar_module: Any = None +validating: Any = None +#: World sizes the matrix runs at (one node): see ``_rank_job.WORLD_SIZES``. +WORLD_SIZES = (2, 4, 8) +#: The ``own_op`` values the cells are driven with: False (the image's op), then True when this target's op is +#: loaded. The last one is the op the target's boundaries call, which the other checks use. +OWN_OPS: Tuple[bool, ...] = (False,) +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +# ---------------------------------------------------------------- inputs + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + """One rank's [m, H] BF16 rows; ``rank`` -1 for tensors every rank holds alike.""" + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return _rows(seed + 7, -1, m, 4.0) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _finalize_inputs(seed: int, rank: int, m: int, topk: int, dropped: int): + """One rank's unfinalized MoE outputs: fc2 [m * topk, H], perm [m, topk], fp32 scales, shared [m, H].""" + g = _gen(seed + 23, rank) + rows = m * topk + fc2 = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + perm = torch.randperm(rows, generator=g, device="cuda").to(torch.int32) + if dropped: + drop = torch.randperm(rows, generator=g, device="cuda")[:dropped] + perm[drop] = -1 + perm = perm.view(m, topk).contiguous() + scales = torch.rand(m, topk, generator=g, device="cuda", dtype=torch.float32) + shared = _rows(seed + 29, rank, m, 0.5) + return fc2, perm, scales, shared + + +# ---------------------------------------------------------------- workspace / producer emulation + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor, exchange_set: int) -> None: + """Store this rank's rows into every rank's exchange, as a producer kernel does before a prepushed call. + + Slot of (token t, cluster CTA c): set * 2 * half + parity(t, c) * half + (t * world + rank) * H / 2 + c * H / 16 + words, parity read from the CTA's flag; -0.0 halves are stored as +0.0 (0x80000000 is the empty word). + """ + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = ( + exchange_set * 2 * half + + parity.view(m, CLUSTER, 1) * half + + (t * WORLD + RANK) * (HIDDEN // 2) + + c * chunk + + w + ) + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +# ---------------------------------------------------------------- gates + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +#: Share of normed elements one BF16 step off the float64 reference: measured at most 1.1e-4 per call over +#: every call of this matrix on GB300 (see ``_assert_normed``). The gate sits ~18x above that and far below what a +#: wrong statistic produces (most elements move). +NORMED_STEP_FRACTION = 0.002 +NORMED_STEP_SEEN: List[float] = [] + + +def _assert_normed(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """At most one BF16 rounding step from the float64 reference, and rarely that. + + The kernel's RMS statistic is a shuffle / warp / cluster tree of FP32 partial sums and its ``rsqrtf`` is + approximate (~2^-22 relative): the scaled value moves by ~1e-6 relative, which changes the BF16 rounding only + for values within that distance of a rounding midpoint -- one step, never two. The fraction bound catches a + statistic that is systematically off (a missing rank, a wrong H), which moves most elements. + """ + assert out.dtype == ref.dtype == torch.bfloat16 and out.shape == ref.shape, where + o, r = out.float(), ref.float() + _, exponent = torch.frexp(torch.maximum(o.abs(), r.abs())) + step = torch.ldexp(torch.ones_like(o), exponent - 8) + diff = (o - r).abs() + assert bool((diff <= step).all()), f"{where}: max diff {diff.max().item()} beyond one BF16 step" + frac = (diff != 0).float().mean().item() + NORMED_STEP_SEEN.append(frac) + assert frac <= NORMED_STEP_FRACTION, f"{where}: {frac:.4f} of normed elements off by one step" + + +def _assert_quant(outs, ref, m: int, where: str) -> None: + """The MXFP8 copy: bit-exact against the recipe applied to the op's own normed rows (a pure function of + them), and equal to the reference's wherever the reference's normed 32-column block is the op's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + mxfp8_rows, + sf_written, + ) + + q, sf = outs[2], outs[3] + own_q, own_sf = mxfp8_rows(outs[0]) + written = sf_written(m, HIDDEN // 32, sf.device) + _assert_bitwise(q, own_q, f"{where} e4m3 rows") + _assert_bitwise(sf[written], own_sf[written], f"{where} UE8M0 scales") + same_block = (outs[0].view(m, -1, 32) == ref[0].view(m, -1, 32)).all(-1) + qb = q.view(torch.uint8).view(m, -1, 32) + rqb = ref[2].view(torch.uint8).view(m, -1, 32) + assert bool((qb == rqb).all(-1)[same_block].all()), ( + f"{where}: e4m3 differs on a block whose normed rows agree" + ) + + +def _check_outputs(outs, ref, m: int, quant: bool, where: str) -> None: + assert len(outs) == (4 if quant else 2), (where, len(outs)) + _assert_bitwise(outs[1], ref[1], f"{where} updated residual") + _assert_normed(outs[0], ref[0], f"{where} normed") + if quant: + _assert_quant(outs, ref, m, where) + + +# ---------------------------------------------------------------- calls + + +def _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, opts, own_op=None): + own_op = OWN_OPS[-1] if own_op is None else own_op + return allreduce_entry.m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + ws.max_tokens, + quant, + prepushed, + opts, + own_op=own_op, + ) + + +def _plain_case(spec: Dict[str, Any], seed: int, where: str, own_op: bool) -> None: + m, maxt, quant, prepushed = spec["m"], spec["maxt"], spec["quant"], spec["prepushed"] + ws = WORKSPACES[maxt] + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] if spec["x2"] else None + residual, norm_w = _residual(seed, m), _norm_w(seed) + x, x2 = xs[RANK], (None if x2s is None else x2s[RANK]) + if prepushed >= 2: + _push(ws, x, 0) + if prepushed == 3: + _push(ws, x2, 1) + with validating(allreduce_entry.m3_mc_allreduce_norm): + outs = _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, spec["opts"], own_op) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, residual, norm_w, EPS, quant, spec.get("drafter", False) + ) + _check_outputs(outs, ref, m, quant, where) + dist.barrier() + + +# ---------------------------------------------------------------- checks + + +def _op_name(own_op: bool) -> str: + return "own op" if own_op else "image op" + + +def check_plain_cells() -> None: + """Every cell of ``m3_mc_allreduce_norm.CELLS`` through ``__call__`` under ``validating``, gated on its + ``reference``, with each op of ``OWN_OPS``.""" + for own_op in OWN_OPS: + for i, cell in enumerate(allreduce_entry.m3_mc_allreduce_norm.CELLS): + _plain_case(cell.spec, 1000 + 17 * i, f"{cell.why} ({_op_name(own_op)})", own_op) + + +def check_finalize_cells() -> None: + """Every cell of ``m3_mc_finalize_allreduce_norm.CELLS``, gated on its ``reference``, with each op of + ``OWN_OPS``.""" + for own_op in OWN_OPS: + _finalize_cells(own_op) + + +def _finalize_cells(own_op: bool) -> None: + op = finalize_entry.m3_mc_finalize_allreduce_norm + for i, cell in enumerate(op.CELLS): + spec = cell.spec + m, maxt, quant, topk = spec["m"], spec["maxt"], spec["quant"], spec["topk"] + seed = 2000 + 17 * i + ws = WORKSPACES[maxt] + per_rank = [_finalize_inputs(seed, r, m, topk, spec["dropped"]) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + with validating(op): + outs = op( + fc2, + perm, + scales, + shared, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + maxt, + quant, + spec["opts"], + own_op=own_op, + ) + torch.cuda.synchronize() + ref = op.reference( + [p[0] for p in per_rank], + [p[1] for p in per_rank], + [p[2] for p in per_rank], + [p[3] for p in per_rank], + residual, + norm_w, + EPS, + quant, + ) + _check_outputs(outs, ref, m, quant, f"{cell.why} ({_op_name(own_op)})") + dist.barrier() + + +def _capture_reference( + updated: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor +) -> None: + """The capture the boundary replaces: m3_spec_tail.capture_norm of the updated residual (no second addend).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_spec_tail, + ) + + m3_spec_tail.capture_norm(updated, None, weight, eps, out) + + +def check_capture_norm_in_boundary() -> None: + """CAPTURE_NORM_IN_BOUNDARY: with cap_out the boundary also writes the capture, BITWISE equal to capture_norm of + its own updated residual, into a strided [M, H] block of a [rows, 3H] capture buffer (other bytes untouched); its + normed / updated / MXFP8 outputs are bitwise the ones of the same call without cap_out. The row kernel (the plain + entry: x, x + x2, prepushed) and the cluster kernel (MoE finalize), 1 / 4 / 8 rows, residual scales from 1e-3 to + 1e3 (the statistic's range), the fc_norm eps 1e-6 and 1e-5; on the op the target's boundaries call.""" + cap_eps_values = (1e-6, 1e-5) + own_op = OWN_OPS[-1] + for case, (m, scale, kind) in enumerate( + ( + (1, 1.0, "plain"), + (4, 1e3, "plain_x2"), + (8, 1e-3, "prepushed"), + (4, 1.0, "prepushed"), + (1, 4.0, "finalize"), + (4, 1e3, "finalize"), + (8, 1e-3, "finalize"), + ) + ): + seed = 5000 + 31 * case + cap_eps = cap_eps_values[case % 2] + ws = WORKSPACES[8] + residual = (_residual(seed, m).float() * scale).to(torch.bfloat16) + norm_w = _norm_w(seed) + cap_w = _norm_w(seed + 101) + quant = case % 3 == 1 + outs_by = [] + for with_cap in (False, True): + buf = torch.full((m + 3, 3 * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + cap = buf[:m, HIDDEN : 2 * HIDDEN] if with_cap else None + if kind == "finalize": + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, m, 4, 1) + fc2 = (fc2.float() * scale).to(torch.bfloat16) + outs = finalize_entry.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), + ws.flags_for(CLUSTER), RANK, WORLD, CLUSTER, False, 8, quant, 7, + cap, cap_w if with_cap else None, cap_eps, own_op=own_op, + ) # fmt: skip + else: + x = (_rows(seed, RANK, m).float() * scale).to(torch.bfloat16) + x2 = ( + (_rows(seed + 3, RANK, m).float() * scale).to(torch.bfloat16) + if kind == "plain_x2" + else None + ) + prepushed = 2 if kind == "prepushed" else 0 + if prepushed: + _push(ws, x, 0) + outs = allreduce_entry.m3_mc_allreduce_norm( + x, x2, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), ws.flags_for(CLUSTER), RANK, WORLD, + CLUSTER, False, 8, quant, prepushed, (7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + cap, cap_w if with_cap else None, cap_eps, own_op=own_op, + ) # fmt: skip + torch.cuda.synchronize() + outs_by.append((outs, buf)) + dist.barrier() + (plain, _), (capped, buf) = outs_by + where = f"capture in boundary, {kind}, {m} rows, scale {scale}" + assert len(plain) == len(capped), where + for i, (a, b) in enumerate(zip(plain, capped)): + if i == 3: + # The scale buffer is torch.empty: compare the bytes the kernel writes (the valid rows'). + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written, + ) + + written = sf_written(m, HIDDEN // 32, a.device) + a, b = a[written], b[written] + _assert_bitwise(b, a, f"{where}: output {i} with vs without cap_out") + ref = torch.empty((m, HIDDEN), dtype=torch.bfloat16, device="cuda") + _capture_reference(capped[1], cap_w, cap_eps, ref) + torch.cuda.synchronize() + _assert_bitwise(buf[:m, HIDDEN : 2 * HIDDEN], ref, f"{where}: capture vs capture_norm") + untouched = buf.clone() + untouched[:m, HIDDEN : 2 * HIDDEN] = 7.0 + assert bool((untouched == 7.0).all()), f"{where}: bytes outside the capture block changed" + dist.barrier() + + +def check_the_target_wrappers_select_bucket_and_options() -> None: + """The target's own wrappers (bucket by token count, option table) on top of the entries, every token count of + the 8-token bucket plus one of each larger bucket, and the finalize wrapper; the MXFP8 publication is the + prequant slot's business and is left off.""" + buckets = ar_module.McBuckets([WORKSPACES[b] for b in BUCKETS]) + for m in (*range(1, 9), 11, 40): + seed = 3000 + m + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + normed, updated = ar_module.m3_mc_allreduce_norm( + xs[RANK], residual, norm_w, EPS, buckets, x2=x2s[RANK] + ) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, False) + _check_outputs([normed, updated], ref, m, False, f"wrapper m={m}") + dist.barrier() + for m in (9, 33): + seed = 3100 + m + per_rank = [_finalize_inputs(seed, r, m, 4, 2) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + normed, updated = ar_module.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, buckets + ) + torch.cuda.synchronize() + ref = finalize_entry.m3_mc_finalize_allreduce_norm.reference( + *[[p[j] for p in per_rank] for j in range(4)], residual, norm_w, EPS, False + ) + _check_outputs([normed, updated], ref, m, False, f"finalize wrapper m={m}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """Two boundaries captured in one graph (8- and 16-token buckets), replayed with fresh inputs. + + The decode graph holds every boundary call; each replay rotates the same per-CTA parities the eager calls do, + so a replay that reused a stale half or skipped a re-arm reads the previous replay's rows and fails the + bitwise residual gate (every replay gets new inputs). + """ + shapes = ((3, 8, 15), (12, 16, 15)) + norm_w = _norm_w(4000) + bufs = [] + for m, maxt, opts in shapes: + bufs.append( + dict( + m=m, + maxt=maxt, + opts=opts, + x=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + x2=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + residual=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + ) + + def body(): + return [ + _call_plain( + WORKSPACES[b["maxt"]], b["x"], b["x2"], b["residual"], norm_w, False, 0, b["opts"] + ) + for b in bufs + ] + + def fill(seed: int): + refs = [] + for j, b in enumerate(bufs): + s = seed + 31 * j + xs = [_rows(s, r, b["m"]) for r in range(WORLD)] + x2s = [_rows(s + 3, r, b["m"]) for r in range(WORLD)] + b["x"].copy_(xs[RANK]) + b["x2"].copy_(x2s[RANK]) + b["residual"].copy_(_residual(s, b["m"])) + refs.append( + allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, b["residual"].clone(), norm_w, EPS, False + ) + ) + torch.cuda.synchronize() + dist.barrier() + return refs + + refs = fill(4100) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm = body() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(warm[j], refs[j], b["m"], False, f"graph warm-up m={b['m']}") + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = body() + for rep in range(4): + refs = fill(4200 + 100 * rep) + graph.replay() + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(captured[j], refs[j], b["m"], False, f"replay {rep} m={b['m']}") + dist.barrier() + del graph + dist.barrier() + + +def check_the_gate_sees_a_missing_rank() -> None: + """The residual gate is not blind to the reduction: dropping any one rank's partial from the reference + moves most elements of the updated residual (each rank's rows are distinct by construction).""" + m, seed = 4, 5000 + xs = [_rows(seed, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + full = allreduce_entry.m3_mc_allreduce_norm.reference(xs, None, residual, norm_w, EPS, False) + for r in range(WORLD): + partial = [x if i != r else torch.zeros_like(x) for i, x in enumerate(xs)] + wrong = allreduce_entry.m3_mc_allreduce_norm.reference( + partial, None, residual, norm_w, EPS, False + ) + frac = (wrong[1] != full[1]).float().mean().item() + assert frac > 0.9, f"dropping rank {r} changes only {frac:.3f} of the residual" + dist.barrier() + + +def _expect_rejected(fn, match: str, where: str) -> None: + try: + fn() + except AssertionError as exc: + assert match in str(exc), f"{where}: rejected for the wrong reason: {exc}" + else: + raise AssertionError(f"{where}: the guard admitted it") + + +def check_guards() -> None: + """``is_valid`` refuses what the op accepts and answers wrongly. Nothing here reaches the kernel: the guard + raises first, on every rank alike, so no rank is left polling.""" + m, seed = 4, 6000 + ws8, ws16 = WORKSPACES[8], WORKSPACES[16] + x, residual, norm_w = _rows(seed, RANK, m), _residual(seed, m), _norm_w(seed) + op = allreduce_entry.m3_mc_allreduce_norm + + def call(ws_uc, flags, norm): + with validating(op): + op( + x, + None, + residual, + norm, + EPS, + ws_uc, + ws8.mc.data_ptr(), + flags, + RANK, + WORLD, + CLUSTER, + False, + 8, + False, + 0, + 15, + ) + + strided = torch.empty(2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[::2] + strided.copy_(norm_w) + _expect_rejected( + lambda: call(ws8.uc, ws8.flags_for(CLUSTER), strided), "contiguous", "strided norm_w" + ) + _expect_rejected( + lambda: call(ws8.uc[: ws8.uc.numel() // 2], ws8.flags_for(CLUSTER), norm_w), + "words", + "short workspace", + ) + _expect_rejected( + lambda: call(ws8.uc, ws16.flags_for(CLUSTER), norm_w), "counters", "another bucket's flags" + ) + + fin = finalize_entry.m3_mc_finalize_allreduce_norm + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, 9, 4, 0) + bad = perm.clone() + bad[0, 0] = fc2.shape[0] + + def call_fin(): + with validating(fin): + fin( + fc2, + bad, + scales, + shared, + _residual(seed, 9), + norm_w, + EPS, + ws16.uc, + ws16.mc.data_ptr(), + ws16.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + 16, + False, + 15, + ) + + _expect_rejected(call_fin, "perm spans", "perm past the FC2 rows") + dist.barrier() + + +CHECKS = ( + check_plain_cells, + check_finalize_cells, + check_capture_norm_in_boundary, + check_the_target_wrappers_select_bucket_and_options, + check_cuda_graph_replays, + check_the_gate_sees_a_missing_rank, + check_guards, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, OWN_OPS, allreduce_entry, finalize_entry, ar_module, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD in WORLD_SIZES, f"the matrix runs at {WORLD_SIZES} ranks on one node, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as _finalize, + ) + + allreduce_entry, finalize_entry, ar_module = _plain, _finalize, _ar + # The same library is loaded on every rank, so the ranks agree on the ops they drive. + OWN_OPS = (False, True) if _ar.own_op() else (False,) + print(f"[rank {RANK}] ops: {', '.join(_op_name(o) for o in OWN_OPS)}", flush=True) + # Spawned by file path, not collected by pytest: reach the sibling helper the plain way (confined to the + # rank processes). /comm -> tests/unittest/_torch/modeling_v2. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peers' polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print( + f"[rank {RANK}] {len(CHECKS)} checks passed; normed one-step fraction max " + f"{max(NORMED_STEP_SEEN):.5f} over {len(NORMED_STEP_SEEN)} calls", + flush=True, + ) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns (one rank each)" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size in WORLD_SIZES, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at {WORLD_SIZES}" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py new file mode 100644 index 000000000000..50e69fc4b290 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py @@ -0,0 +1,403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification of the MoE engine's all-reduce push (``m3_moe_engine`` with ``arp``), 8 ranks, two trays (SM107). + +The decode forward's MoE boundary (MOE_PUSH rows): the engine pushes its bf16 sum (routed + shared) into set A of +the boundary's multicast exchange while it writes it, and the boundary all-reduce runs with ``prepushed`` 2, reading +the pushed words instead of loading and pushing the rows itself. Certified against the unpushed boundary the target +runs otherwise: the same engine call without ``arp``, then ``m3_mc_allreduce_norm`` with ``prepushed`` 0. Every +output must be bit for bit the same: the engine's own output tensor, the updated residual, the normed rows and, with +``quant``, the MXFP8 hand-off it publishes. Both read the same bf16 words, so the reduction and the norm see equal +inputs; any difference is a push that stored the wrong word, slot or parity. + +Each rank holds its own expert shard (operand seed per rank) and the same hidden rows, residual and norm weight, as +the TP2 decode step does. Eager calls alternate pushed and unpushed in both orders so both call parities of the +exchange are exercised; a CUDA graph of the pushed engine + boundary is replayed on fresh rows against eager +unpushed references. + +Its own launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is: run plainly it spawns one rank process per device +in CUDA_VISIBLE_DEVICES and kills them all on the first failure or on the deadline (a rank whose peer never pushes +polls forever). The collected entry point is ``test_modeling_v2_m3_moe_push_op_matrix.py``. + + CUDA_VISIBLE_DEVICES=0,1 python _m3_moe_push_op_matrix.py +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_moe_engine requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1200 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +MAX_TOKENS = 8 +EMPTY_WORD = -2147483648 +RESERVED_SMS = 24 +OPERAND_SEED = 20261013 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +ar = None +sf_written: Any = None +engine_op: Any = None +ar_push_desc: Any = None +moe_state: Any = None +operands: Any = None +prequant: Any = None +RANK = 0 +WORLD = 1 +WS: Any = None + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return (torch.randn(m, HIDDEN, generator=_gen(seed + 7), device="cuda") * 4.0).to( + torch.bfloat16 + ) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _make_workspace(): + """The 8-token bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed + group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=4) + words = 4 * MAX_TOKENS * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=MAX_TOKENS, + ) + ws.flags_for(CLUSTER) + return ws + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +_STATES = {} +_LAYERS = {} + + +def _layer(m: int, case: str): + """This rank's engine layer for ``m`` tokens: its own expert shard; ``concentrated`` routes every token to the + same four experts.""" + if m not in _STATES: + sms = torch.cuda.get_device_properties(0).multi_processor_count + _STATES[m] = moe_state( + torch.device("cuda"), num_tokens=m, num_ctas=sms - RESERVED_SMS, shared=True + ) + key = (m, case) + if key not in _LAYERS: + ops = operands.experts(OPERAND_SEED + RANK) + bias = ops.bias.clone() + if case == "concentrated": + bias[list(operands.CONCENTRATED)] += 1.0 + _LAYERS[key] = _STATES[m].layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=ops.shared, + ) + return _LAYERS[key] + + +def _desc(): + flags = WS.flags_for(CLUSTER) + return ar_push_desc( + WS.mc.data_ptr(), flags.data_ptr(), RANK, WORLD, WS.max_tokens, flags.device + ) + + +def _boundary(x, residual, norm_w, quant, prepushed): + prequant.SLOT.clear() + outs = ar.m3_mc_allreduce_norm( + x, residual, norm_w, EPS, WS, use_pdl=True, quant=quant, prepushed=prepushed + ) + assert outs is not None, "the boundary all-reduce declined its operands" + pub = prequant.SLOT.pending + prequant.SLOT.clear() + return outs, (None if pub is None else (pub[3].clone(), pub[4].clone())) + + +def _unpushed(hidden, layer, residual, norm_w, quant): + out = engine_op(hidden, layer) + (normed, res), pub = _boundary(out, residual.clone(), norm_w, quant, 0) + return out.clone(), normed.clone(), res.clone(), pub + + +def _pushed(hidden, layer, residual, norm_w, quant): + out = engine_op(hidden, layer, arp=_desc()) + (normed, res), pub = _boundary(out, residual.clone(), norm_w, quant, 2) + return out.clone(), normed.clone(), res.clone(), pub + + +def _compare(got, ref, quant: bool, where: str) -> None: + _assert_bitwise(got[0], ref[0], f"{where} engine output") + _assert_bitwise(got[2], ref[2], f"{where} updated residual") + _assert_bitwise(got[1], ref[1], f"{where} normed") + if quant: + assert got[3] is not None and ref[3] is not None, f"{where}: no MXFP8 hand-off published" + _assert_bitwise(got[3][0], ref[3][0], f"{where} e4m3 rows") + # Only the scale atoms of the m rows are written; the swizzled buffer's padding is not. + written = sf_written(got[0].shape[0], HIDDEN // 32, got[3][1].device) + _assert_bitwise(got[3][1][written], ref[3][1][written], f"{where} UE8M0 scales") + + +# ---------------------------------------------------------------- checks + + +def check_pushed_boundary_is_the_unpushed_one() -> None: + """Every token count, both routing cases, with and without the MXFP8 hand-off, both call orders.""" + for m in range(1, MAX_TOKENS + 1): + for case in ("random", "concentrated"): + layer = _layer(m, case) + for quant in (False, True): + for seed in (1, 2): + s = ( + 1000 * m + + 10 * seed + + (5 if quant else 0) + + (500 if case == "concentrated" else 0) + ) + hidden = operands.hidden_states(m, s) + residual, norm_w = _residual(s, m), _norm_w(s) + where = f"m={m} {case} quant={quant} seed={seed}" + if seed == 1: + ref = _unpushed(hidden, layer, residual, norm_w, quant) + got = _pushed(hidden, layer, residual, norm_w, quant) + else: + got = _pushed(hidden, layer, residual, norm_w, quant) + ref = _unpushed(hidden, layer, residual, norm_w, quant) + torch.cuda.synchronize() + _compare(got, ref, quant, where) + + +def check_graph_replays() -> None: + """A captured pushed engine + prepushed boundary, replayed on fresh rows, equals the eager unpushed path.""" + for m in (1, 4, 8): + layer = _layer(m, "random") + hidden = operands.hidden_states(m, 90000 + m).clone() + residual = _residual(90000 + m, m).clone() + norm_w = _norm_w(90000 + m) + desc = _desc() + for _ in range( + 2 + ): # warm both parities eagerly; every push is consumed by its prepushed boundary + _pushed(hidden, layer, residual, norm_w, False) + _unpushed(hidden, layer, residual, norm_w, False) + torch.cuda.synchronize() + dist.barrier() + res_in = residual.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + g_out = engine_op(hidden, layer, arp=desc) + g_normed, g_res = ar.m3_mc_allreduce_norm( + g_out, res_in, norm_w, EPS, WS, use_pdl=True, prepushed=2 + ) + torch.cuda.synchronize() + dist.barrier() + for rep in range(6): + s = 91000 + 100 * m + rep + fresh_h, fresh_r = operands.hidden_states(m, s), _residual(s, m) + hidden.copy_(fresh_h) + res_in.copy_(fresh_r) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + got = (g_out.clone(), g_normed.clone(), g_res.clone(), None) + ref = _unpushed(fresh_h, layer, fresh_r, norm_w, False) + torch.cuda.synchronize() + _compare(got, ref, False, f"graph m={m} replay {rep}") + dist.barrier() + + +CHECKS = [check_pushed_boundary_is_the_unpushed_one, check_graph_replays] + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, ar, sf_written, engine_op, ar_push_desc, moe_state, operands, prequant, WS + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + local = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", local) + ) + RANK, WORLD = rank, world + assert WORLD == 8, f"the TP8 push is certified at 8 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_prequant as _pq # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written as _sf_written, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + M3MoeState, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + ar_push_desc as _desc_fn, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as _op, + ) + + ar, prequant, moe_state, ar_push_desc, engine_op = _ar, _pq, M3MoeState, _desc_fn, _op + sf_written = _sf_written + # Spawned by file path, not collected by pytest: the sibling moe/ operands the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "moe")) + import _m3_moe_operands as _operands + + operands = _operands + WS = _make_workspace() + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peer's polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + ) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +_SRUN_FLAG = "--srun-rank" + + +def _srun_rank() -> int: + """One rank per Slurm task over two trays (srun -N2 --ntasks-per-node 4): RANK / WORLD_SIZE / LOCAL_RANK from + Slurm; the rendezvous is MASTER_ADDR / MASTER_PORT exported by the job (the first node).""" + os.environ["RANK"] = os.environ["SLURM_PROCID"] + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"] + return _run_one_rank() + + +if __name__ == "__main__": + if _SRUN_FLAG in sys.argv: + sys.exit(_srun_rank()) + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_rank_job.py new file mode 100644 index 000000000000..dedb967cd195 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_rank_job.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run this target's collective rank job from pytest, and report what it did. + +The same contract as ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``: the +launcher beside this file owns the per-device rank processes and the deadline +that kills a wedged collective; this only selects the devices, starts it in a +fresh interpreter (by file path: this tree is not a package) and turns its exit +code into an assertion. + +World size. The tp8 target runs 8 ranks over two GB300 trays of 4 GPUs (one +NVLink domain); a single-node unit-test job cannot form that group. The +collective entries take the world size as a runtime argument and their cells' +specs do not name one (token counts, buckets, push modes, epilogue options), so +the jobs here certify every cell at each world size in ``WORLD_SIZES`` the +node's visible devices can form: 2 and 4 on a GB300 tray, 8 only on a node with +8 devices. A world size the node cannot form is skipped with that reason, never +emulated; the 8-rank, two-tray group itself is exercised only by the target's +multi-node runs. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +TARGET_WORLD_SIZE = 8 +"""tp8's world size: 8 ranks over two GB300 trays, which one node's job cannot form.""" + +WORLD_SIZES = (2, 4, TARGET_WORLD_SIZE) +"""The world sizes the jobs certify the (world-independent) cells at, where the node has that many devices.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the launcher's own deadline, so its message wins the race.""" + + +def visible_devices() -> list: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + return [d for d in visible.split(",") if d.strip()] + import torch + + return [str(i) for i in range(torch.cuda.device_count())] + + +def require_world(world_size: int) -> str: + """The first ``world_size`` visible devices; skips (with the reason) when the node has fewer.""" + devices = visible_devices() + if len(devices) < world_size: + where = ( + "the tp8 target's own 8-rank group spans two trays and is exercised only by its multi-node runs" + if world_size == TARGET_WORLD_SIZE + else "run this job on a node with more devices" + ) + pytest.skip( + f"{world_size} ranks need {world_size} devices; {len(devices)} visible on this node ({where})" + ) + return ",".join(devices[:world_size]) + + +def _load_launcher_constants(launcher: Path): + """Read the launcher's module-level budget without importing it by name (it imports only torch at + module scope).""" + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str, world_size: int) -> str: + """Run ``__op_matrix``'s launcher over ``world_size`` devices (skipped when the node has fewer); returns + its stdout.""" + assert world_size in WORLD_SIZES, world_size + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=require_world(world_size)) + + # The ranks import tensorrt_llm absolutely. From a source checkout that is not installed, put the repo root + # (tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root) on their path; from a tests + # tree run against an installed package, they import the installed one. + repo_root = Path(__file__).resolve().parents[8] + if (repo_root / "tensorrt_llm").is_dir(): + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p + ) + else: + assert importlib.util.find_spec("tensorrt_llm") is not None, ( + f"no tensorrt_llm/ at the repo root {repo_root} and no installed tensorrt_llm" + ) + + constants = _load_launcher_constants(launcher) + timeout = constants.DEADLINE_S + _LAUNCHER_GRACE_S + + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode} at world size {world_size}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) + return completed.stdout diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_allreduce.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_allreduce.py new file mode 100644 index 000000000000..c19a7461a368 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_allreduce.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the ``m3_allreduce`` certification matrix. + +The matrix is ``_m3_allreduce_op_matrix.py`` beside this file, its own launcher; see ``_rank_job``. It drives every +cell of the entry against its reference, then requires the fused draft loop's boundaries built on it (plain +all-reduce, then the fused add + RMSNorm) to equal stock Eagle3's bit for bit under every single-node strategy, and +reports both against an fp64 reference as a warning, one line per strategy and case. + +The tp8 drafter runs at 8 ranks over two trays; a single-node job cannot form that group, so the matrix runs at each +world size of ``_rank_job.WORLD_SIZES`` the node's devices can form (the claim, draft boundary = stock boundary, does +not depend on the rank count) and skips the others. MNNVL, the strategy the two-tray group may take, is multi-node +only and not driven here. +""" + +import warnings + +import _rank_job +import pytest +import torch + +# Per world size, _rank_job.require_world skips the sizes this node's devices cannot form. +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +_RESULT_TAG = "[ulp]" + + +# Spawns one rank process per device of the world size; under xdist workers would fight for them. +@pytest.mark.no_xdist +@pytest.mark.parametrize("world_size", _rank_job.WORLD_SIZES, ids=lambda w: f"world{w}") +def test_m3_allreduce_op_matrix(world_size) -> None: + out = _rank_job.run("m3_allreduce", world_size) + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + f"draft boundary (= stock Eagle3) vs fp64 at {world_size} ranks, max BF16 ulps\n" + + "\n".join(lines), + stacklevel=1, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..0a9a7d39c177 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the multicast all-reduce + RMSNorm entries' certification matrix. + +Covers both ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm`` (one +kernel, one exchange protocol). The matrix is ``_m3_mc_allreduce_norm_op_matrix.py`` +beside this file, its own launcher; it drives every cell of each entry's +``CELLS`` against the entry's ``reference``. See ``_rank_job`` for why. + +The tp8 target runs these entries at 8 ranks over two trays, which a single-node +job cannot form. The world size is a runtime argument of the kernel and no cell's +spec names one, so the matrix certifies every cell at each world size of +``_rank_job.WORLD_SIZES`` the node's devices can form (2 and 4 on a 4-device tray, +8 on an 8-device node) and skips the others; it does not emulate the two-tray group. +""" + +import _rank_job +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + finalize_partial, +) + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + + +# Spawns one rank process per device of the world size; under xdist workers would fight for them. +@pytest.mark.no_xdist +@pytest.mark.parametrize("world_size", _rank_job.WORLD_SIZES, ids=lambda w: f"world{w}") +def test_m3_mc_allreduce_norm_op_matrix(world_size) -> None: + _rank_job.run("m3_mc_allreduce_norm", world_size) + + +def test_the_cells_are_world_independent() -> None: + """What lets a single-node job certify the tp8 entries: no cell's spec names a world size.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm, + m3_mc_finalize_allreduce_norm, + ) + + for op in ( + m3_mc_allreduce_norm.m3_mc_allreduce_norm, + m3_mc_finalize_allreduce_norm.m3_mc_finalize_allreduce_norm, + ): + assert op.CELLS + for cell in op.CELLS: + assert not {"world", "world_size", "rank"} & set(cell.spec), cell.why + + +def test_finalize_partial_rounds_the_routed_sum_before_the_shared_expert() -> None: + """As stock: the finalized routed output is BF16, and the shared expert is added to it (shared.add_(routed)). + Four routed rows of 1.0 at FP32 routing weight 0.2509765625 sum to 1.00390625, which rounds to BF16 1.0; plus the + shared -1 that is 0, where adding before the rounding would leave 0.00390625.""" + fc2 = torch.ones((4, 8), dtype=torch.bfloat16) + perm = torch.arange(4, dtype=torch.int32).view(1, 4) + scales = torch.full((1, 4), 0.2509765625, dtype=torch.float32) + shared = torch.full((1, 8), -1.0, dtype=torch.bfloat16) + out = finalize_partial(fc2, perm, scales, shared) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(out, shared.clone().add_(routed)) + assert torch.equal(out, torch.zeros(1, 8, dtype=torch.bfloat16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..546ad3e11ef7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_embedding entry of the MiniMax-M3 tp8 target: an exact gather with zero rows outside the +table.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.embedding.m3_embedding import ( # noqa: E501 + m3_embedding as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" + +VOCAB, HIDDEN = 4096, 6144 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + ids = torch.randint( + 0, VOCAB, (spec["rows"],), generator=g, device="cuda", dtype=spec["id_dtype"] + ) + if spec["pad"]: + ids[1::3] = -1 + with validating(op): + out = op(ids, table) + assert torch.equal(out, op.reference(ids, table)) + if spec["pad"]: + assert not out[1::3].any(), "padded ids must embed to zero rows" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py new file mode 100644 index 000000000000..e44ad56d743d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" + +import torch + + +def swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): + """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" + v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) + return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..509c1c09cebd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_bf16_linear entry of the MiniMax-M3 tp8 target (the drafter's projections).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_bf16_linear import ( # noqa: E501 + m3_bf16_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the + row, RMS well under one; measured on GB300: worst 0.5 unit.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w = (torch.randn(spec["n"], spec["k"], generator=g, device="cuda") * 0.02).to(torch.bfloat16) + with validating(op): + out = op(x, w) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(x, w), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py new file mode 100644 index 000000000000..67796dc0302f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise oracle of the dense MLP's swigluoai C++ op (ops/gemm/m3_dense_swiglu, +trtllm::minimax_m3_sm107_tp8_swiglu_oai) against the Triton silu_and_mul kernel the dense MLP runs today +(modules/swiglu.py at limit 7, alpha 1.702, beta 1): every bf16 gate bit pattern against special and random up values, +every up bit pattern against special and random gates, random pairs, and the dense MLP's shapes at every decode row +count, eager and replayed from a CUDA graph, with and without the programmatic launch.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.silu_and_mul +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_dense_swiglu import ( # noqa: E501 + ALPHA, + BETA, + LIMIT, + m3_dense_swiglu, +) + +assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_swiglu_oai"): + pytest.skip( + "needs the TP8 op library with minimax_m3_sm107_tp8_swiglu_oai", allow_module_level=True + ) + +# The TP8 dense MLP's activation width (gate_up 3072 -> d 1536). +DENSE_D = modeling.DENSE_DOWN[1] + + +def _triton(x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=LIMIT, swiglu_alpha=ALPHA, swiglu_beta=BETA + ) + + +def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" + return out + + +def _row(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + return torch.cat([gate.view(1, -1), up.view(1, -1)], dim=1).contiguous() + + +def _all_bf16() -> torch.Tensor: + return ( + torch.arange(1 << 16, dtype=torch.int32, device="cuda").to(torch.int16).view(torch.bfloat16) + ) + + +def _specials(n: int, seed: int) -> list: + g = torch.Generator(device="cuda").manual_seed(seed) + vals = [0.0, -0.0, 1.0, -1.0, 6.99, 7.0, 7.01, -7.0, -7.01, 1e-38, -1e-38, 1e-40, 88.0, -88.0, 3e38, -3e38, + float("inf"), float("-inf"), float("nan")] # fmt: skip + rows = [torch.full((n,), v, dtype=torch.bfloat16, device="cuda") for v in vals] + for scale in (0.05, 1.0, 4.0, 30.0): + rows.append((scale * torch.randn(n, generator=g, device="cuda")).to(torch.bfloat16)) + rows.append( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + return rows + + +def _assert_bitwise(got: torch.Tensor, ref: torch.Tensor, what: str) -> None: + bad = (got.view(torch.int16) != ref.view(torch.int16)).nonzero().flatten() + assert bad.numel() == 0, f"{what}: {bad.numel()} mismatches, first at {bad[:4].tolist()}" + + +def test_every_gate_pattern_is_bitwise_the_triton_kernel() -> None: + gate = _all_bf16() + for j, up in enumerate(_specials(gate.numel(), 1)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"up set {j}") + + +def test_every_up_pattern_is_bitwise_the_triton_kernel() -> None: + up = _all_bf16() + for j, gate in enumerate(_specials(up.numel(), 2)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"gate set {j}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: + g = torch.Generator(device="cuda").manual_seed(100 + seed) + n = 1 << 22 + + def pattern() -> torch.Tensor: + bits = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + return bits.to(torch.int16).view(torch.bfloat16) + + x = _row(pattern(), pattern()) + _assert_bitwise(_op(x), _triton(x), f"seed {seed}") + + +@pytest.mark.parametrize("pdl", [True, False], ids=["pdl", "plain"]) +@pytest.mark.parametrize("rows", range(1, modeling.DECODE_MAX_TOKENS + 1)) +def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pdl: bool) -> None: + """The TP8 dense MLP's gate_up rows (a GEMV output's range plus values past the clamps), eager and replayed from a + CUDA graph after fresh rows are written in place, against the Triton kernel on the same rows.""" + g = torch.Generator(device="cuda").manual_seed(2000 + rows) + x = (3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, DENSE_D : DENSE_D + 8] = x[:, :8] + _assert_bitwise(_op(x, pdl), _triton(x), f"eager rows {rows}") + out = torch.empty(rows, DENSE_D, dtype=torch.bfloat16, device="cuda") + _op(x, pdl) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out.copy_(_op(x, pdl)) + x.copy_((3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(out, _triton(x), f"graph rows {rows}") + + +def test_the_op_declines_other_parameters_and_layouts() -> None: + x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + + +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: + """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton + kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" + rows, d = cell.spec["rows"], cell.spec["d"] + g = torch.Generator(device="cuda").manual_seed(3000 + rows) + x = (3.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, d : d + 8] = x[:, :8] + _assert_bitwise(_op(x), _triton(x), cell.why) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py new file mode 100644 index 000000000000..e63d7f3f11ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_draft_lowm entry: the Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M +kernels (direct SIMT and split-K tensor-core), at the TP8 drafter shards, at the rows its table certifies. + +Every cell is held to the fp64 reference within the derived band. The drafter's projection call +(``M3Eagle3Draft.linear``) at a covered row count is the entry's call bit for bit, at an uncovered one the in-tree +GEMM's. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + TACTICS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + m3_draft_lowm as op, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" + +# Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). +# +# Kernel and reference multiply the same BF16 operands (every product exact in fp32); the kernel sums in fp32 (the +# direct kernel one dot product in K order, the split-K kernel tensor-core blocks and then its split partials), the +# reference in fp64, and both round once to BF16. With w ~ N(0, 1 / K) and x ~ N(0, 1) the outputs are O(1) and the +# fp32 order difference is ~sqrt(K) * 2^-24 of the summed magnitudes, a few 1e-6 at K = 18432: far below half a BF16 +# ulp, so the two roundings disagree by at most one ulp of the element -- up to two units of the row scale for an +# element in its row maximum's binade -- and only where the sum sits that close to a rounding boundary, a few in ten +# thousand elements, hence 0.1 ulp RMS. The same derivation as the decode GEMV's band (test_modeling_v2_m3_mx_gemv), +# with no slack under it: a second ulp would mean a wrong sum. +ELEMENT_ULP = 2.0 +RMS_ULP = 0.1 + +#: The drafter's projections (N, K) and their weight keys on ``M3Eagle3Draft.w``. +DRAFTER_LINEARS = { + (1280, 12288): "qkv", + (6144, 1024): "o", + (4608, 6144): "gu", + (6144, 2304): "dn", + (6144, 18432): "fc", +} + + +@pytest.fixture(autouse=True) +def _policy(): + modeling.set_require(False) + modeling.reset_logged() + yield + modeling.set_require(False) + modeling.reset_logged() + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + return x, w + + +def test_the_table_is_the_tp8_drafter() -> None: + assert set(TACTICS) == set(DRAFTER_LINEARS) + assert {c.spec["nk"] for c in op.CELLS} == set(TACTICS) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + kind, tactic = tactic_for(n, k, m) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + args = ( + x, + w, + out, + kind, + tactic, + True, + ) # the target launches with PDL when the environment enables it + certifying(op) + assert op(*args) is out + ref = op.reference(*args) + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=RMS_ULP) + # PDL only changes the launch attribute. + plain = torch.empty_like(out) + op(x, w, plain, kind, tactic, False) + assert torch.equal(plain.view(torch.int16), out.view(torch.int16)) + + +def _draft(linears=DRAFTER_LINEARS): + """A drafter holding only its projection weights and the low-M bookkeeping (the module is not built).""" + draft = object.__new__(modeling.M3Eagle3Draft) + draft.w = {key: build(*nk, 1, seed=nk[0] + nk[1])[1] for nk, key in linears.items()} + draft._lowm_built, draft._lowm_engaged = set(), set() + draft._push_arm, draft._pushed = None, 0 + return draft + + +@pytest.fixture(scope="module") +def prebuilt(): + draft = _draft() + assert draft.prebuild_linears() == len(op.CELLS) + assert draft.prebuild_linears() == 0, "a second prebuild finds nothing left to compile" + return draft + + +@pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) +def test_the_projection_is_the_entry_and_falls_back_past_its_rows(prebuilt, nk) -> None: + n, k = nk + key = DRAFTER_LINEARS[nk] + w = prebuilt.w[key] + use_pdl = get_env_enable_pdl() + modeling.set_require(True) + for m in ROWS[nk]: + assert modeling.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt.linear(x, key) + kind, tactic = tactic_for(n, k, m) + want = op(x, w, torch.empty(m, n, dtype=x.dtype, device="cuda"), kind, tactic, use_pdl) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)), m + # Rows the table does not certify (the b2 rows, and past it): the in-tree GEMM, silently, even under require. + for m in sorted(set(range(1, modeling.DECODE_MAX_TOKENS + 2)) - set(ROWS[nk])): + assert not modeling.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt.linear(x, key) + assert torch.equal(got.view(torch.int16), F.linear(x, w).view(torch.int16)), m + + +def test_a_projection_outside_the_table_is_the_in_tree_gemm() -> None: + draft = _draft({(256, 512): "o"}) + x, _ = build(256, 512, 1, seed=3) + assert draft._lowm_linear(draft.w["o"], x) is None + got = draft.linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_a_covered_call_not_prebuilt_raises_under_require() -> None: + draft = _draft({(6144, 1024): "o"}) + x, _ = build(6144, 1024, 1, seed=5) + modeling.set_require(True) + with pytest.raises(modeling.FastPathUnavailable, match="not compiled before warmup"): + draft.linear(x, "o") + # Outside require it falls back to the in-tree GEMM. + modeling.set_require(False) + got = draft.linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_guard_refuses_a_strided_activation() -> None: + x, w = build(6144, 1024, 2, seed=6) + wide = torch.zeros(2, 2048, dtype=torch.bfloat16, device="cuda") + wide[:, :1024] = x + kind, tactic = tactic_for(6144, 1024, 2) + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op( + wide[:, :1024], + w, + torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda"), + kind, + tactic, + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..149d230a929c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp32_linear entry of the MiniMax-M3 tp8 target (the router beyond the decode routers' rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """torch's default fp32 matmul precision on this hardware may be TF32 (~2^-11 relative per product): the band is + stated in fp32 ulps of each row's largest logit and covers it (K = 6144 TF32 products: ~2^13 ulps worst element); + measured on GB300 with the engine's defaults: worst element ~1700 ulps, RMS ~300.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(spec["e"], spec["k"], generator=g, device="cuda") * 0.02 + with validating(op): + out = op(x, w) + assert out.dtype == torch.float32 and out.shape == (spec["m"], spec["e"]) + assert_within_ulp(out, op.reference(x, w), element_ulp=8192.0, rms_ulp=2048.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..918aa9e5ffab --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv entry: the decode MXFP8 GEMV on a producer-quantized activation. + +Operands as the target holds them: the weight through mxfp8_quantize with the swizzled scale layout (the MXFP8 linear +method's storage), the activation through mxfp8_quantize with the swizzled layout (what the boundary all-reduce and +the attention merge publish). +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" + + +# Element band, in ULP units of the row scale: see compare. +ELEMENT_ULP = 1.5 + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Element ELEMENT_ULP of the row scale, relative RMS 0.1 ulp. + + Kernel and reference multiply bit-identical dequantized operands; the kernel sums in fp32 (tensor-core blocks, + then the split-K partials), the reference in fp64. Over K <= 6144 that order difference is ~1e-6 relative, so the + two bf16 roundings disagree only where the sum sits within that of a rounding boundary: one bf16 ulp of that + element, which for an element in its row maximum's binade is up to two of ``assert_within_ulp``'s units of the row + scale, and a few in ten thousand elements in aggregate. Any fp32 sum shows it: torch's IEEE fp32 matmul of the + same dequantized operands measured at most 1.22 units against the fp64 reference over every row count 1..48 on + the o_proj and the QKV (6 seeds each, GB300), the kernel at most 1.25, both at most 0.009 ulp RMS. The band + leaves the margin between that and a full extra ulp. The kernel is the same on SM107 (the tp2 table only picks + split and ring), and its TP2 projections' K is 1536..6144. + """ + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=0.1) + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + return x, xq.view(m, k), xsf, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + _, xq, xsf, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(xq, xsf, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + compare(out, op.reference(xq, xsf, w, wsf)) + + +def test_an_unserved_shape_returns_none() -> None: + """Shapes outside the tuned table fall back at the call site; the op says so rather than computing.""" + _, xq, xsf, w, wsf = build(256, 512, 1, seed=1) + assert op(xq, xsf, w, wsf) is None + + +def test_guard_refuses_linear_activation_scales() -> None: + """A linear [M, K / 32] scale buffer is read as the swizzled one.""" + x, _, _, w, wsf = build(6144, 4096, 4, seed=2) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, False) + with pytest.raises(AssertionError, match="swizzled"): + with validating(op): + op(xq.view(4, 4096), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + _, xq, xsf, w, wsf = build(6144, 4096, 1, seed=3) + wide = torch.zeros(6144, 8192, dtype=w.dtype, device="cuda") + wide[:, :4096] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :4096], wsf) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..c8c9461b44da --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv_bf16 entry: the decode MXFP8 GEMV quantizing a bf16 activation in its prologue. + +The weight as the MXFP8 linear method stores it (mxfp8_quantize, swizzled scales). Besides the reference band, the +prologue's quantization is held to mxfp8_quantize itself: the GEMV on a bf16 x must equal, bit for bit, the +prequantized GEMV on mxfp8_quantize(x) -- the claim the target relies on when it mixes the two paths. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv_bf16 import ( # noqa: E501 + m3_mx_gemv_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" + +# The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. +ELEMENT_ULP = 1.5 + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return x, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(x, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + # The band of test_modeling_v2_m3_mx_gemv.compare: the quantized operands are the reference's exactly (checked + # below), the sums differ only in order. + assert_within_ulp(out, op.reference(x, w, wsf), element_ulp=ELEMENT_ULP, rms_ulp=0.1) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + pre = m3_mx_gemv(xq.view(m, k), xsf, w, wsf, use_pdl=True) + assert pre is not None and torch.equal(out.view(torch.int16), pre.view(torch.int16)), ( + "the prologue's MXFP8 quantization differs from mxfp8_quantize" + ) + + +def test_a_strided_activation_is_declined() -> None: + """The prologue reads x densely; the op declines a strided x (None) rather than computing on it.""" + x, w, wsf = build(6144, 4096, 2, seed=1) + wide = torch.zeros(2, 8192, dtype=torch.bfloat16, device="cuda") + wide[:, :4096] = x + assert op(wide[:, :4096], w, wsf) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py new file mode 100644 index 000000000000..d1793092cf8b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Oracle for #117, the sparse QKV + index GEMV with the producer in its owner epilogue (mx_gemv +``m3_mx_gemv_producer``): q, index-q and every cache byte (NVFP4 K / V data and scales, FP8 index-K) bit for bit equal +to today's chain, the GEMV storing the packed BF16 row and then the producer kernel on it. Both producer builds are +references: the library's ``minimax_m3_producer_pw`` and, when the target's extension is loaded, its own +``minimax_m3_sm107_tp8_producer_pw`` (the copy the target dispatches, with its pre-wait step-input reads).""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.mx_gemv import ( # noqa: E501 + op as gemv, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +N, K = 12 * 128, 6144 +HEADS_Q = 8 +PAGE = 128 +PAGES = 512 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _case(tokens, seed, scale=1.0, bad_slots=False): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(N, K, generator=g, device="cuda") * (scale / K**0.5)).to(torch.bfloat16) + x = torch.randn(tokens, K, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + norms = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots and tokens >= 3: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return xq.view(tokens, K), xsf, wq.view(N, K), wsf, slots, quant, norms, rotary, positions + + +def _caches(): + return ( + torch.zeros(PAGES, 2, 1, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, 1, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _chain(case, caches, producer): + xq, xsf, wq, wsf, slots, quant, norms, rotary, positions = case + packed = gemv.m3_mx_gemv(xq, xsf, wq, wsf, use_pdl=True) + assert packed is not None, "the sparse QKV shape must be served" + if producer == "lib": + q, iq = torch.ops.trtllm.minimax_m3_producer_pw( + packed, *caches, slots, quant, EPS, *norms, rotary, positions, HEADS_Q, True, True + ) + else: + q, iq = torch.ops.trtllm.minimax_m3_sm107_tp8_producer_pw( + packed, *caches, slots, quant, EPS, *norms, rotary, positions, True, True, True + ) + torch.cuda.synchronize() + return q, iq + + +def _fused(case, caches): + xq, xsf, wq, wsf, slots, quant, norms, rotary, positions = case + tokens = xq.shape[0] + q = torch.zeros(tokens, HEADS_Q, 128, dtype=torch.float8_e4m3fn, device="cuda") + iq = torch.zeros(tokens, 1, 128, dtype=torch.float8_e4m3fn, device="cuda") + kv_data, kv_scale, index_k = caches + ok = gemv.m3_mx_gemv_producer( + xq, xsf, wq, wsf, q, iq, kv_data, kv_scale, index_k, slots, quant, EPS, *norms, rotary, positions, + use_pdl=True, + ) # fmt: skip + assert ok, "the producer GEMV must serve the sparse QKV shape at <= 8 tokens" + torch.cuda.synchronize() + return q, iq + + +def _producers(): + out = ["lib"] + if hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_producer_pw"): + out.append("tp8") + return out + + +@pytest.mark.parametrize( + "tokens, seed, scale, bad_slots", + [ + (1, 21, 1.0, False), + (2, 22, 1.0, False), + (4, 23, 1.0, False), + (4, 24, 40.0, False), + (8, 25, 1.0, False), + (8, 26, 0.01, False), + (8, 27, 1.0, True), + ], + ids=["t1", "t2", "t4", "t4-large", "t8", "t8-tiny", "t8-bad-slots"], +) +@pytest.mark.parametrize("producer", _producers()) +def test_producer_gemv_is_bitwise_the_chain(tokens, seed, scale, bad_slots, producer) -> None: + case = _case(tokens, seed, scale, bad_slots) + ref_caches = _caches() + rq, riq = _chain(case, ref_caches, producer) + caches = _caches() + q, iq = _fused(case, caches) + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, a in zip(names, (rq, riq, *ref_caches)): + assert bool(a.contiguous().view(torch.uint8).any()), f"reference {name} is all zero" + for name, a, b in zip(names, (rq, riq, *ref_caches), (q, iq, *caches)): + a8, b8 = a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8) + diff = (a8 != b8).nonzero() + assert diff.numel() == 0, ( + f"{name}: {diff.shape[0]} bytes differ, first at {tuple(diff[0].tolist())}" + ) + + +def test_producer_gemv_is_deterministic_and_graph_replays() -> None: + """Three eager runs and a captured replay write the same bytes.""" + case = _case(4, 31) + outs = [] + for _ in range(3): + caches = _caches() + q, iq = _fused(case, caches) + outs.append([t.contiguous().view(torch.uint8).clone() for t in (q, iq, *caches)]) + xq, xsf, wq, wsf, slots, quant, norms, rotary, positions = case + q = torch.zeros(4, HEADS_Q, 128, dtype=torch.float8_e4m3fn, device="cuda") + iq = torch.zeros(4, 1, 128, dtype=torch.float8_e4m3fn, device="cuda") + caches = _caches() + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + gemv.m3_mx_gemv_producer( + xq, xsf, wq, wsf, q, iq, *caches, slots, quant, EPS, *norms, rotary, positions, use_pdl=True + ) # fmt: skip + graph.replay() + torch.cuda.synchronize() + outs.append([t.contiguous().view(torch.uint8).clone() for t in (q, iq, *caches)]) + for run in outs[1:]: + for a, b in zip(outs[0], run): + assert torch.equal(a, b) + + +def test_declines_outside_its_shape() -> None: + """More than 8 tokens, or a projection other than the 12-head packed one, launches nothing.""" + case = _case(8, 41) + xq, xsf, wq, wsf, slots, quant, norms, rotary, positions = case + wide = torch.cat([xq, xq], dim=0) + q = torch.zeros(16, HEADS_Q, 128, dtype=torch.float8_e4m3fn, device="cuda") + iq = torch.zeros(16, 1, 128, dtype=torch.float8_e4m3fn, device="cuda") + xsf16 = torch.ops.trtllm.mxfp8_quantize( + torch.zeros(16, K, dtype=torch.bfloat16, device="cuda"), True + )[1] + caches = _caches() + assert not gemv.m3_mx_gemv_producer( + wide, xsf16, wq, wsf, q, iq, *caches, slots, quant, EPS, *norms, rotary, positions + ) # fmt: skip diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..4374a72aa8a0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_gemm entry of the MiniMax-M3 tp8 target (the decode-graph GEMM on a prequantized +activation).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most + 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, False), + element_ulp=2.0, + rms_ulp=0.25, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..dcff36cd9b93 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_linear entry of the MiniMax-M3 tp8 target: MXFP8 quantize + GEMM on each backend the +engine selects (native CUTLASS, FlashInfer, the per-bucket tuned quantize + GEMM).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" + +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and + one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" + spec = cell.spec + if ( + spec["path"] == "flashinfer" + and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None + ): + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + args = (x, w, w_sf, alpha, torch.bfloat16, *_FLAGS[spec["path"]]) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py new file mode 100644 index 000000000000..457d56dd8698 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for MX-TAC: the pinned decode-graph MXFP8 decision (CuTe DSL quantize + CuTe DSL GEMM) gives the same bits +as every other candidate stock's AutoTuner can pick (TRT-LLM quantize, CUTLASS GEMM), the same on every call.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.modeling import ( # noqa: E501 + _pinned_mxfp8_gemm, + _pinned_mxfp8_quantize, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + +assert torch.cuda.is_available(), "MX-TAC requires a CUDA device" + +# The TP8 per-rank projections (ops/gemm/m3_mx_gemv) at decode rows the GEMVs do not serve (9..16) and a few beyond. +CASES = [ + ((1536, 6144), 16), + ((1280, 6144), 12), + ((6144, 1024), 9), + ((768, 6144), 16), + ((6144, 384), 12), + ((3072, 6144), 24), + ((6144, 1536), 32), + ((1536, 6144), 48), +] + + +@pytest.mark.parametrize("nk,m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_pinned_decision_matches_every_stock_candidate(nk, m, scale) -> None: + n, k = nk + g = torch.Generator(device="cuda").manual_seed(n + k + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = (torch.randn(m, k, generator=g, device="cuda") * scale).to(torch.bfloat16) + x[0, ::32] *= 1e3 # one outlier per 32-element block on the first row + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + wq = wq.view(n, k) + xq, xsf = _pinned_mxfp8_quantize(x) + tq, tsf = torch.ops.trtllm.mxfp8_quantize(x, True) + assert torch.equal(xq.view(torch.uint8).view(m, k), tq.view(torch.uint8).view(m, k)) + assert torch.equal(xsf.view(torch.uint8).flatten(), tsf.view(torch.uint8).flatten()) + xq = xq.view(m, k) + out = _pinned_mxfp8_gemm(xq, xsf, wq, wsf, torch.bfloat16) + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + for tactic in (FlashInferMXFP8GemmRunner.CUTE_DSL, FlashInferMXFP8GemmRunner.CUTLASS): + ref = runner.forward([xq, xsf, wq, wsf], tactic=tactic) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), tactic + again = _pinned_mxfp8_gemm(xq, xsf, wq, wsf, torch.bfloat16) + assert torch.equal(out.view(torch.int16), again.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py new file mode 100644 index 000000000000..1d1ffee82db5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise oracle of the mx_gemv swigluoai prologue arithmetic (kernels/mx_gemv/_swiglu_oai_bf16x2) against the +Triton silu_and_mul kernel the dense MLP runs today (modules/swiglu.py, alpha 1.702, limit 7, beta 1): every bf16 gate +bit pattern against special and random up values, and every up bit pattern against special and random gates.""" + +import cutlass +import cutlass.cute as cute +import pytest +import torch +from cutlass.cute.runtime import from_dlpack + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.mx_gemv.m3_mx_gemv_kernel import ( # noqa: E501 + SWIGLU_OAI_ALPHA, + SWIGLU_OAI_LIMIT, + _swiglu_oai_bf16x2, +) + +assert torch.cuda.is_available(), "the swigluoai prologue oracle requires a CUDA device" + +THREADS = 256 + + +@cute.kernel +def _swiglu_words_kernel(gate: cute.Tensor, up: cute.Tensor, out: cute.Tensor, n: cutlass.Int32): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + i = bidx * cutlass.Int32(THREADS) + tidx + if i < n: + out[i] = _swiglu_oai_bf16x2(gate[i], up[i]) + + +@cute.jit +def _swiglu_words(gate: cute.Tensor, up: cute.Tensor, out: cute.Tensor, n: cutlass.Int32): + _swiglu_words_kernel(gate, up, out, n).launch( + grid=((n + THREADS - 1) // THREADS, 1, 1), block=(THREADS, 1, 1) + ) + + +_compiled = {} + + +def _prologue(gate_bf16: torch.Tensor, up_bf16: torch.Tensor) -> torch.Tensor: + """The prologue helper over flat bf16 tensors of even length (pairs packed into int32 words).""" + g = gate_bf16.contiguous().view(torch.int32) + u = up_bf16.contiguous().view(torch.int32) + out = torch.empty_like(g) + args = [from_dlpack(t, assumed_align=4) for t in (g, u, out)] + fn = _compiled.get("f") + if fn is None: + fn = _compiled["f"] = cute.compile(_swiglu_words, *args, cutlass.Int32(g.numel())) + fn(*args, cutlass.Int32(g.numel())) + return out.view(torch.bfloat16) + + +def _triton(gate_bf16: torch.Tensor, up_bf16: torch.Tensor) -> torch.Tensor: + """The dense MLP's activation today: torch.ops.trtllm.silu_and_mul on one [1, 2d] row per call site layout.""" + x = torch.cat([gate_bf16.view(1, -1), up_bf16.view(1, -1)], dim=1) + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=SWIGLU_OAI_LIMIT, swiglu_alpha=SWIGLU_OAI_ALPHA, swiglu_beta=1.0 + ).view(-1) + + +def _all_bf16() -> torch.Tensor: + return ( + torch.arange(1 << 16, dtype=torch.int32, device="cuda").to(torch.int16).view(torch.bfloat16) + ) + + +def _specials(n: int, seed: int) -> list: + g = torch.Generator(device="cuda").manual_seed(seed) + vals = [0.0, -0.0, 1.0, -1.0, 6.99, 7.0, 7.01, -7.0, -7.01, 1e-38, -1e-38, 1e-40, 88.0, -88.0, 3e38, -3e38, + float("inf"), float("-inf"), float("nan")] # fmt: skip + rows = [torch.full((n,), v, dtype=torch.bfloat16, device="cuda") for v in vals] + for scale in (0.05, 1.0, 4.0, 30.0): + rows.append((scale * torch.randn(n, generator=g, device="cuda")).to(torch.bfloat16)) + rows.append( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + return rows + + +def _assert_bitwise(got: torch.Tensor, ref: torch.Tensor, what: str) -> None: + bad = (got.view(torch.int16) != ref.view(torch.int16)).nonzero().flatten() + assert bad.numel() == 0, f"{what}: {bad.numel()} mismatches, first at {bad[:4].tolist()}" + + +def test_every_gate_pattern_is_bitwise_the_triton_kernel() -> None: + gate = _all_bf16() + for j, up in enumerate(_specials(gate.numel(), 1)): + _assert_bitwise(_prologue(gate, up), _triton(gate, up), f"up set {j}") + + +def test_every_up_pattern_is_bitwise_the_triton_kernel() -> None: + up = _all_bf16() + for j, gate in enumerate(_specials(up.numel(), 2)): + _assert_bitwise(_prologue(gate, up), _triton(gate, up), f"gate set {j}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: + g = torch.Generator(device="cuda").manual_seed(100 + seed) + n = 1 << 22 + gate = ( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + up = ( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + _assert_bitwise(_prologue(gate, up), _triton(gate, up), f"seed {seed}") + + +# The GEMV with the activation in its prologue against today's two launches (Triton silu_and_mul, then the bf16 GEMV): +# bit for bit on the dense MLP's down projection at every decode row count. +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E402 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.mx_gemv import ( # noqa: E402, E501 + op as gemv, +) + + +@pytest.mark.parametrize("m", range(1, modeling.DECODE_MAX_TOKENS + 1)) +def test_down_gemv_with_the_activation_prologue_is_bitwise_two_launches(m: int) -> None: + n, k = modeling.DENSE_DOWN + g = torch.Generator(device="cuda").manual_seed(1000 + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + wq = wq.view(n, k) + # A gate_up GEMV output's range, plus values past the clamps. + gate_up = (3.0 * torch.randn(m, 2 * k, generator=g, device="cuda")).to(torch.bfloat16) + gate_up[:, :8] = torch.tensor( + [8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16 + ) + act = torch.ops.trtllm.silu_and_mul( + gate_up, swiglu_limit=SWIGLU_OAI_LIMIT, swiglu_alpha=SWIGLU_OAI_ALPHA, swiglu_beta=1.0 + ) + ref = gemv.m3_mx_gemv_bf16(act.contiguous(), wq, wsf, use_pdl=True) + out = gemv.m3_mx_gemv_swiglu(gate_up, wq, wsf, use_pdl=True) + assert ref is not None and out is not None, ( + "the dense down projection must be served at every decode row count" + ) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py new file mode 100644 index 000000000000..f207d8896427 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rule 1b gate a for the TP8 decode GEMV order override (kernels/mx_gemv/op.STRIDED_K_TILES: strided k-tiles): +for every overridden (projection, rows) the output is the same on 3 eager calls and on a CUDA graph replay +(deterministic), and its max abs / max rel error vs the fp64 product of the dequantized operands stays within the +shipped order's band (the shared table's config on the same operands, over several seeds).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.mx_gemv import ( # noqa: E501 + op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + mxfp8_dequant, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +CASES = [ + (nk, m) for nk in (QKV_SPARSE, QKV, O_PROJ, DENSE_GATE_UP, DENSE_DOWN) for m in (1, 2, 4, 8) +] + + +def _err(y, ref): + d = (y.double() - ref).abs() + return float(d.max()), float((d / ref.abs().amax(dim=1, keepdim=True).clamp_min(1e-30)).max()) + + +@pytest.mark.parametrize("nk, m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +def test_gate_a(nk, m) -> None: + n, k = nk + cfg = op.get_config(n, k, m) + assert cfg is not None and len(cfg) > 2 and cfg[2] == 0, ( + "the override is off or does not cover this call" + ) + shipped = cfg[:2] # the contiguous order + worst_new, worst_ship = [0.0, 0.0], [0.0, 0.0] + for seed in range(4): + g = torch.Generator(device="cuda").manual_seed(1000 * seed + m) + w = (torch.randn(n, k, device="cuda", generator=g) * 2).to(torch.float8_e4m3fn) + ws = torch.randint(118, 124, (n * k // 32,), dtype=torch.uint8, device="cuda", generator=g) + x = (torch.randn(m, k, device="cuda", generator=g) * (0.5 + seed)).to(torch.bfloat16) + q, s = torch.ops.trtllm.mxfp8_quantize(x, True) + ref = ( + mxfp8_dequant(q.view(torch.uint8), s).double() + @ mxfp8_dequant(w.view(torch.uint8), ws).double().t() + ) + ys = [m3_mx_gemv(q, s, w, ws, True, cfg, True, None) for _ in range(3)] + assert all(y is not None for y in ys), "the GEMV declined the override" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + out = [] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out.append(m3_mx_gemv(q, s, w, ws, True, cfg, True, None)) + graph.replay() + torch.cuda.synchronize() + for y in ys[1:] + [out[0]]: + assert torch.equal(ys[0].view(torch.int16), y.view(torch.int16)), ( + "not deterministic (eager / graph)" + ) + ship = m3_mx_gemv(q, s, w, ws, True, shipped, True, None) + torch.cuda.synchronize() + for acc, y in ((worst_new, ys[0]), (worst_ship, ship)): + e = _err(y, ref) + acc[0], acc[1] = max(acc[0], e[0]), max(acc[1], e[1]) + print( + f"\nR1B1 tp8 {n}x{k} m={m} cfg {cfg}: max abs / rel vs fp64 {worst_new[0]:.3e} / {worst_new[1]:.3e} | shipped " + f"order {worst_ship[0]:.3e} / {worst_ship[1]:.3e}", + flush=True, + ) + # The band: the shipped order's own error (bf16 output rounding dominates), with 25 % headroom. + assert worst_new[0] <= 1.25 * worst_ship[0] and worst_new[1] <= 1.25 * worst_ship[1] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..b3b710ec185f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_moe_operands.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One tp8 rank's MiniMax-M3 MoE layer operands for the decode-engine entries' tests. + +The routed experts are random checkpoint-format NVFP4 tensors ([w3 (up); w1 (gate)] rows, e4m3 block scales) put +through TRT-LLM's own TRTLLM-Gen NVFP4 loader shuffle (``NVFP4TRTLLMGenFusedMoEMethod._shuffle_all_experts``), which +writes the buffers the engines read in place. The per-expert scalars follow the loader's conventions (act_beta and +act_clamp divided by fc31_alpha, fc31_scale_c = fc2_input_scale x fc31_alpha) and are sized so the FC1 sums land +around the SwiGLU-OAI clamp and the outputs near unit scale. The shared expert's MXFP8 weights come from +mxfp8_quantize with the swizzled scale layout, as the MXFP8 linear method stores them. +""" + +import functools +from types import SimpleNamespace + +import torch + +# Hidden size, the routed experts' intermediate per rank at TP8 (1536 / 8 x 2: every expert on every rank), experts, +# top-k, and the shared expert's intermediate per rank. +H, I_TP, NUM_EXPERTS, TOP_K = 6144, 384, 128, 4 +SHARED_I = 384 +SWIGLU_ALPHA, SWIGLU_LIMIT = 1.702, 7.0 +ROUTED_SCALING = 2.0 +MINIMAX2_ROUTING = 5 # RoutingMethodType.MiniMax2 +# NVFP4 global scales: 448 x 6 over the amax each side is sized for. +FC31_INPUT_SCALE = 448.0 * 6.0 / 8.0 +FC2_INPUT_SCALE = 448.0 * 6.0 / 24.0 + + +def _rand_e4m3_scales(shape, gen) -> torch.Tensor: + """E4M3 block scales in [0.5, 2): biased exponents 6 and 7, any mantissa.""" + return torch.randint(0x30, 0x40, shape, dtype=torch.uint8, device="cuda", generator=gen).view( + torch.float8_e4m3fn + ) + + +@functools.lru_cache(maxsize=None) +def experts(seed: int = 20261013): + """The routed experts and the router of one MoE layer, in the layout and conventions of the TRTLLM-Gen loader. + + Returns a SimpleNamespace: ``tensors`` (the ten expert operands in ``M3MoeState.layer``'s order), ``bias``, + ``router_weight``, ``fc31_input_scale`` and ``shared`` (gate_up, its scales, down, its scales). + """ + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + gen = torch.Generator(device="cuda").manual_seed(seed) + e = NUM_EXPERTS + w31 = torch.randint( + 0, 256, (e, 2 * I_TP, H // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + w2 = torch.randint(0, 256, (e, H, I_TP // 2), dtype=torch.uint8, device="cuda", generator=gen) + w31s = _rand_e4m3_scales((e, 2 * I_TP, H // 16), gen) + w2s = _rand_e4m3_scales((e, H, I_TP // 16), gen) + module = SimpleNamespace( + w3_w1_weight=w31, + w2_weight=w2, + w3_w1_weight_scale=w31s, + w2_weight_scale=w2s, + is_gated_activation=True, + bias=False, + ) + NVFP4TRTLLMGenFusedMoEMethod()._shuffle_all_experts(module) + torch.cuda.synchronize() + + # Raw FC1 sums: rms(x) x FC31_INPUT_SCALE x rms(E2M1 value) x rms(block scale) x sqrt(H) ~ 1.0e5 at rms(x) = 1; + # fc31_alpha brings them to ~1.5 in the dequantized domain, the FC2 alpha the outputs to ~1. + base = 1.5 / (FC31_INPUT_SCALE * 2.93 * 1.3 * H**0.5) + fc31_alpha = base * (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) + fc2_alpha = (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) / ( + FC2_INPUT_SCALE * 2.93 * 1.3 * I_TP**0.5 + ) + tensors = ( + w31, + w31s, + w2, + w2s, + fc31_alpha, + FC2_INPUT_SCALE * fc31_alpha, # fc31_scale_c + torch.full((e,), SWIGLU_ALPHA, device="cuda"), # act_alpha + 1.0 / fc31_alpha, # act_beta: the (up + 1) offset over fc31_alpha + SWIGLU_LIMIT / fc31_alpha, # act_clamp over fc31_alpha + fc2_alpha, + ) + bias = (0.02 * torch.randn(e, device="cuda", generator=gen)).float() + router_weight = (torch.randn(e, H, device="cuda", generator=gen) / H**0.5).float().contiguous() + fc31_input_scale = torch.tensor([FC31_INPUT_SCALE], dtype=torch.float32, device="cuda") + + gu = (0.02 * torch.randn(2 * SHARED_I, H, device="cuda", generator=gen)).to(torch.bfloat16) + dn = (0.05 * torch.randn(H, SHARED_I, device="cuda", generator=gen)).to(torch.bfloat16) + gu_q, gu_sf = torch.ops.trtllm.mxfp8_quantize(gu, True) + dn_q, dn_sf = torch.ops.trtllm.mxfp8_quantize(dn, True) + shared = (gu_q.view(2 * SHARED_I, H), gu_sf, dn_q.view(H, SHARED_I), dn_sf) + return SimpleNamespace( + tensors=tensors, + bias=bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared=shared, + ) + + +# Routing cases: random (up to 4 M distinct experts), and every token on the same four experts (the bias outweighs +# any sigmoid difference), which leaves later slot groups empty. +CONCENTRATED = (5, 40, 77, 126) + + +def layer_for(state, case: str, shared: bool = True): + """One layer on ``state`` over :func:`experts` for a routing case. ``shared=False``: the shared expert's weights + are zero, so the engine returns the routed part alone (bf16(0 + bf16(routed))).""" + ops = experts() + bias = ops.bias.clone() + if case == "concentrated": + bias[list(CONCENTRATED)] += 1.0 + shared_weights = ops.shared + if not shared: + gu, gu_sf, dn, dn_sf = ops.shared + shared_weights = (torch.zeros_like(gu), gu_sf, torch.zeros_like(dn), dn_sf) + return state.layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=shared_weights, + ) + + +def stock_routed( + hidden: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """The stock TRTLLM-Gen NVFP4 MoE (fp4_quantize, then fp4_block_scale_moe_runner with SwiGLU-OAI) on the given + routing, over :func:`experts`: what the built-in MiniMaxM3MoE computes for the routed part.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + + ops = experts() + w31, w31s, w2, w2s, fc31_alpha, scale_c, act_alpha, act_beta, act_clamp, fc2_alpha = ops.tensors + xq, xsf = torch.ops.trtllm.fp4_quantize(hidden, ops.fc31_input_scale, 16, False, False) + out = torch.ops.trtllm.fp4_block_scale_moe_runner( + None, None, xq, xsf.view(torch.float8_e4m3fn).flatten(), + w31, w31s, None, act_alpha, act_beta, act_clamp, w2, w2s, None, + scale_c, fc31_alpha, fc2_alpha, + NUM_EXPERTS, TOP_K, None, None, I_TP, 0, NUM_EXPERTS, ROUTED_SCALING, MINIMAX2_ROUTING, True, 0, + topk_weights=topk_weights.to(torch.bfloat16), topk_ids=topk_ids.to(torch.int32), + ) # fmt: skip + return out[0] if isinstance(out, (list, tuple)) else out + + +def hidden_states(m: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(m, H, device="cuda", generator=gen).to(torch.bfloat16) + + +def check_routed_part(op, state, m: int, seed: int) -> None: + """Exact, independent of the reference: the engine's routing (its route_dbg record: ids, then the bf16 weights' + fp32 bits) is the reference routing bit for bit, and its routed part is the stock TRTLLM-Gen MoE's output on that + routing bit for bit -- the claim the engines are built on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=False) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + stock = stock_routed(hidden, ids, weights) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)), ( + "routed part differs from TRTLLM-Gen's" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py new file mode 100644 index 000000000000..7a88ec1adca9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared operands and gates of the router entries' tests (m3_router_gemv, m3_router_gemv_fp4, m3_router_tc: one +kernel source).""" + +import zlib + +import torch + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv import ( # noqa: E501 + m3_router_gemv as gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv import ( # noqa: E501 + near_tie, + nvfp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 + m3_router_gemv_fp4 as gemv_fp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__all__ = [ + "SF_SCALE", + "_build", + "_compare_fp4", + "_compare_logits", + "_seed", + "gemv", + "gemv_fp4", + "tc", +] + +#: The experts' NVFP4 input scale the tc cells use: 448 * 6 / 8, block scales well inside E4M3. +SF_SCALE = 336.0 + + +def _build(m: int, e: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(e, k, generator=g, device="cuda") * 0.02 + return x, w + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +def _compare_logits(out: torch.Tensor, ref: torch.Tensor) -> None: + """fp32 logits that differ from a true fp32 product only in summation order. + + Over K = 6144 products a reordered fp32 sum moves by a few ulps of the sum of magnitudes, which near + cancellation is ~sqrt(K) times the row's largest logit's ulp at most. Gated in fp32 ulps of each row's largest + logit: 64 element-wise, 8 in relative RMS. Measured on GB300: GEMV 1.6 / 1.4 ulp, tensor-core router 3.8 / 2.0 + ulp at 2..48 tokens; the tensor-core logits against a true fp32 product (the TF32 rounding the band must see) + sit at 2200+ / 1700 ulp. + """ + assert out.dtype == ref.dtype == torch.float32 and out.shape == ref.shape + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=8.0) + + +def _codes(packed: torch.Tensor) -> torch.Tensor: + """[M, K / 2] packed e2m1 bytes -> [M, K] 4-bit codes, low nibble first.""" + lo, hi = packed & 0xF, packed >> 4 + return torch.stack((lo, hi), dim=-1).reshape(packed.shape[0], -1) + + +def _compare_fp4( + q: torch.Tensor, sf: torch.Tensor, x: torch.Tensor, sf_scale: torch.Tensor +) -> None: + """Scales bit-exact; codes bit-exact except within ``TIE_WINDOW`` of an e2m1 midpoint, where the kernel's + approximate reciprocals choose the side (and then only by one code).""" + ref_q, ref_sf, scaled = nvfp4_quantize(x, sf_scale) + assert q.dtype == torch.uint8 and q.shape == ref_q.shape + assert torch.equal(sf, ref_sf), f"{(sf != ref_sf).sum().item()} block scales differ" + got, want = _codes(q), _codes(ref_q) + differ = got != want + assert not bool((differ & ~near_tie(scaled)).any()), ( + f"{(differ & ~near_tie(scaled)).sum().item()} e2m1 codes differ away from a rounding tie" + ) + magnitude = ((got & 7).int() - (want & 7).int()).abs() + assert bool((magnitude[differ] <= 1).all()), "a tie resolved more than one code away" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..8836a2e7726d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe_engine entry: the 1..8-token MoE engine on its caller-owned state. + +One Rubin (SM107) rank of the tp8 deployment: 128 experts (all local, routed and shared intermediate 384 per +rank), top 4, hidden 6144, the shared expert folded in, ``num_ctas`` = SMs - 24 as the target builds it. Operands: +``_m3_moe_operands``. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + HIDDEN, + INTER, + SHARED_INTER, + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import ( # noqa: E402 + I_TP, + SHARED_I, + H, + check_routed_part, + hidden_states, + layer_for, +) +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +_STATES = {} + + +def test_the_operands_are_the_tp8_engine_geometry() -> None: + assert (H, I_TP, SHARED_I) == (HIDDEN, INTER, SHARED_INTER) == (6144, 384, 384) + + +def test_the_engine_refuses_the_tp4_width() -> None: + with pytest.raises(ValueError, match="384"): + M3MoeState(torch.device("cuda"), num_tokens=1, inter=768) + + +def _state(m: int) -> M3MoeState: + if m not in _STATES: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + _STATES[m] = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + return _STATES[m] + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """The engine against the plain-torch MoE: BAND_ELEMENT bf16 ulp of the row's largest magnitude per element, + BAND_RMS ulp relative RMS. + + Both sides quantize the same operands with the same recipes (the activation's and the intermediate's NVFP4 codes + and block scales, the shared expert's MXFP8); they differ in summation order (tensor-core fp32 against fp64) and in + the PTX approximations (rcp / ex2 / tanh). Such a difference changes the output only where it moves a value across + a rounding boundary: one bf16 ulp of a slot output or of the sum, or, rarely, one intermediate NVFP4 code, which + reaches the output as one weight column's share. Measured over every cell, both routing cases and 3 seeds each + on the tp4 engine (GB300, intermediate 768): at most 1.48 ulp element, 0.18 ulp RMS; the tp8 engine runs the + same recipes over half the intermediate. The band leaves ~3x on both; a misread recipe (the activation's + 1 / (sf x (1 / sc)) spelling of the reciprocal on ~0.06% of its codes) measured 12 ulp element, 7 ulp RMS. + """ + assert_within_ulp(out, ref, element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + # The workspace is tagged per call and never reset: a repeated call is the same call. + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + compare(out, op.reference(hidden, layer)) + + +@pytest.mark.parametrize("m", (1, 2, 8)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +def test_layers_on_one_state_do_not_leak_into_each_other() -> None: + """Two layers share one state's workspace, in stream order: A, B, A gives A's output twice, bit for bit, and the + caller's ``out`` receives the same rows as a fresh tensor.""" + state = _state(4) + a, b = layer_for(state, "random"), layer_for(state, "concentrated") + hidden = hidden_states(4, seed=7) + first = op(hidden, a) + op(hidden, b) + out = torch.empty(4, H, dtype=torch.bfloat16, device="cuda") + second = op(hidden, a, out=out) + assert second.data_ptr() == out.data_ptr() + assert torch.equal(first.view(torch.int16), out.view(torch.int16)) + + +@pytest.mark.parametrize( + "bad", + ("misaligned hidden", "fp32 out", "short out", "long out", "int32 arp"), +) +def test_malformed_buffers_are_rejected_by_the_call(bad) -> None: + """Why the entry has no is_valid: the compiled call checks each buffer against its signature and raises.""" + state = _state(2) + layer = layer_for(state, "random") + hidden, kw = hidden_states(2, seed=3), {} + if bad == "misaligned hidden": + hidden = torch.randn(2 * H + 8, device="cuda").to(torch.bfloat16)[2 : 2 + 2 * H].view(2, H) + elif bad == "fp32 out": + kw["out"] = torch.empty(2, H, dtype=torch.float32, device="cuda") + elif bad == "short out": + kw["out"] = torch.empty(1, H, dtype=torch.bfloat16, device="cuda") + elif bad == "long out": + kw["out"] = torch.empty(3, H, dtype=torch.bfloat16, device="cuda") + else: + kw["arp"] = torch.zeros(8, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + op(hidden, layer, **kw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..ca6b047d862d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_tc entry of the MiniMax-M3 tp8 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + + +@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +def test_tc_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) + sf_scale = torch.tensor([SF_SCALE], dtype=torch.float32, device="cuda") if spec["fp4"] else None + with validating(tc): + outs = tc(x, w, sf_scale, False) + ref = tc.reference(x, w, sf_scale, False) + assert len(outs) == len(ref) == (3 if spec["fp4"] else 1) + _compare_logits(outs[0], ref[0]) + if spec["fp4"]: + _compare_fp4(outs[1], outs[2], x, sf_scale) + + +def test_tc_reference_is_tf32_and_the_gate_sees_it() -> None: + """The tensor-core router rounds the weight to TF32: against a true fp32 product its logits sit far outside + the band, so the TF32 rounding in ``reference`` is load-bearing rather than decoration.""" + x, w = _build(16, 128, 6144, 5) + out = tc(x, w, None, False)[0] + with pytest.raises(AssertionError): + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +def test_guard_refuses_a_misaligned_activation() -> None: + """A row view starting 8 bytes into the buffer keeps a legal row stride but not a legal bulk-copy source.""" + x, w = _build(2, 128, 6144, 3) + flat = torch.empty(2 * 6144 + 4, dtype=torch.bfloat16, device="cuda") + shifted = flat[4:].view(2, 6144) + shifted.copy_(x) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(tc): + tc(shifted, w, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..05f9e80edbd9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/m3_embedding_firstnorm entry: the decode embedding gather + the first layer's RMSNorm. + +The table is the model's full bf16 [200064, 6144] embedding shape (random values at the checkpoint's scale), eps the +checkpoint's 1e-6. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + m3_embedding_firstnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" + +H = 6144 +EPS = 1e-6 + + +@functools.lru_cache(maxsize=None) +def _operands(): + g = torch.Generator(device="cuda").manual_seed(20261013) + table = (0.02 * torch.randn(VOCAB, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return table, weight + + +def _buffers(rows): + return (torch.empty(rows, H, dtype=torch.bfloat16, device="cuda"), + torch.empty(rows, H, dtype=torch.bfloat16, device="cuda")) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + rows, id_dtype = cell.spec["rows"], cell.spec["id_dtype"] + table, weight = _operands() + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device="cuda").to(id_dtype) + raw, out = _buffers(rows) + with validating(op): + normed, residual = op(ids, table, weight, EPS, raw, out) + assert normed.data_ptr() == out.data_ptr() and residual.data_ptr() == raw.data_ptr() + ref_normed, ref_raw = op.reference(ids, table, weight, EPS, raw, out) + # The gather is a copy: exact. + assert torch.equal(raw.view(torch.int16), ref_raw.view(torch.int16)) + # The norm: the kernel's fp32 sum of squares (warp butterflies, then across 4 warps) and fast rsqrt against an + # fp64 reference, then x * rstd * w rounded once to bf16. The fp32 side is ~1e-6 relative off, so the two roundings + # disagree only next to a bf16 rounding boundary: one ulp at most per element, a small fraction in aggregate. + # Measured over these cells (3 seeds each, GB300): at most 0.51 ulp element, 0.026 ulp RMS. + assert_within_ulp(out, ref_normed, element_ulp=1.0, rms_ulp=0.1) + + +def test_ids_outside_the_table_give_zero_rows() -> None: + """Documented, not refused: an id outside [0, V) gathers a zero row whose norm is zero.""" + table, weight = _operands() + ids = torch.tensor([-1, 5, VOCAB], dtype=torch.int64, device="cuda") + raw, out = _buffers(3) + op(ids, table, weight, EPS, raw, out) + assert not raw[0].any() and not raw[2].any() and not out[0].any() and not out[2].any() + assert torch.equal(raw[1], table[5]) + + +@pytest.mark.parametrize("rows", range(1, 9)) +def test_against_the_stock_embedding_and_rmsnorm(rows) -> None: + """What this target certifies on SM107 against the built-in path it replaces (the model's Embedding gather, then + the first layer's RMSNorm module), including a constant row, large and tiny rows and both vocab ends: the gather + bit for bit, the norm within one bf16 ulp (the sum-of-squares tree differs from flashinfer's: measured ~5e-6 of + the elements one ulp off over 12000 random rows).""" + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + table, weight = _operands() + table = table.clone() + g = torch.Generator(device="cuda").manual_seed(rows) + table[1] = 3.0 + table[2] = (300.0 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + table[3] = (1e-4 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + norm = RMSNorm(hidden_size=H, eps=EPS, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + norm.weight.copy_(weight) + for ids in ( + torch.randint(0, VOCAB, (rows,), generator=g, device="cuda"), + torch.tensor([1, 2, 3, 0, VOCAB - 1, 1, 2, 3][:rows], device="cuda"), + ): + raw, out = _buffers(rows) + normed, residual = op(ids, table, weight, EPS, raw, out) + x = torch.nn.functional.embedding(ids, table) + assert torch.equal(residual.view(torch.int16), x.view(torch.int16)) + stock = norm(x) + off = normed != stock + assert int(off.sum()) <= 4, f"{int(off.sum())} normed elements differ from stock" + assert_within_ulp(normed, stock, element_ulp=1.0, rms_ulp=0.05) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..15f7afe40867 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs shared by the tp4 target's spec entry tests: TP4 dimensions, single-device stand-ins for the cross-rank +exchanges, and the comparison bands.""" + +from __future__ import annotations + +import zlib + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) + +#: MiniMax-M3 at TP8: hidden size, vocabulary, the target head's shard and RMSNorm epsilon. +HIDDEN = 6144 +VOCAB = 200064 +WORLD = 8 +TARGET_SHARD = VOCAB // WORLD +EPS = 1e-6 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): + """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" + uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") + keys = torch.zeros(slot_rows, dtype=torch.int64, device="cuda") + counters = torch.tensor([0, call], dtype=torch.int32, device="cuda") + return uc, keys, counters + + +def plant_peers( + uc: torch.Tensor, + counters: torch.Tensor, + own_best: torch.Tensor, + rank: int, + world: int, + shard: int, + g: torch.Generator, + slot_rows: int = SHARD_ROWS, + ties: bool = False, +) -> None: + """Write the other ranks' keys for each row into this rank's slots, as their multicast pushes would. + + Each peer's maximum is the own row maximum shifted by +-0.25 (well clear of the accumulation-order band), so + some rows are won by a peer and some by this rank; with ``ties`` a peer on an even row reports exactly the own + maximum, and the lower global index must win. + """ + if world == 1: + return + m = own_best.shape[0] + half = int(counters[1].item()) & 1 + for r in range(world): + if r == rank: + continue + sign = torch.where(torch.rand(m, generator=g, device="cuda") < 0.5, -1.0, 1.0) + value = own_best.float() + 0.25 * sign + if ties: + value = torch.where(torch.arange(m, device="cuda") % 2 == 0, own_best.float(), value) + index = r * shard + torch.randint(0, shard, (m,), generator=g, device="cuda") + keys = make_keys(value, index) + for t in range(m): + uc[(half * slot_rows + t) * world + r] = keys[t] + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..605c96c773cc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_capture_norm entry (Eagle3 capture with fc_norm folded in).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_capture_norm import ( # noqa: E501 + m3_capture_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" + +LAYERS = 3 + + +def _case(spec, seed): + g = gen(seed) + m, slot = spec["m"], spec["slot"] + h = randn((m, HIDDEN), g, 2.0) + r = randn((m, HIDDEN), g, 2.0) if spec["residual"] else None + w = (1.0 + randn((HIDDEN,), g, 0.1).float()).to(torch.bfloat16) + buf = torch.full((m, LAYERS * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + dst = buf[:, slot * HIDDEN : (slot + 1) * HIDDEN] + return (h, r, w, EPS, dst), buf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer + untouched.""" + args, buf = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + slot = cell.spec["slot"] + others = torch.cat([buf[:, : slot * HIDDEN], buf[:, (slot + 1) * HIDDEN :]], dim=1) + assert bool((others == 7.0).all()) + + +def test_guard_refuses_a_strided_weight() -> None: + (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + w2 = torch.stack([w, w], dim=1)[:, 0] + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(h, r, w2, eps, dst) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py new file mode 100644 index 000000000000..28d5cbefc927 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU test: a first-draft-inputs reader that skips the drafter's fc_norm (the capture norm is armed) gets the captured +rows unchanged when every captured layer was written with the step's norm, and otherwise a copy with only the raw +layers normalized by their own fc_norm (M3Eagle3Worker._m3_renorm_raw_captures); the per-forward reset clears the +record and the unconsumed fc fork result.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) + +M3Eagle3Worker = modeling.M3Eagle3Worker + +H = 4 +LAYERS = (1, 29, 56) + + +def _worker(normed): + w = object.__new__(M3Eagle3Worker) + w._m3_capture_normed = set(normed) + return w + + +@pytest.fixture(autouse=True) +def _scaling_norm(monkeypatch): + """The RMSNorm the worker applies, as ``x * w`` (CPU): enough to tell which blocks were normalized here.""" + monkeypatch.setattr(modeling, "flashinfer_rmsnorm", lambda x, w, eps: x * w) + modeling.set_require(False) + + +def _drafter(): + # fc_norm i scales by (i + 2). + return SimpleNamespace(w={f"fc_norm{i}": i + 2 for i in range(len(LAYERS))}, rms_eps=1e-6) + + +def test_all_captured_layers_normed_returns_the_rows_as_they_lie() -> None: + hs = torch.randn(4, H * len(LAYERS)) + out = _worker(set(LAYERS))._m3_renorm_raw_captures( + hs, SimpleNamespace(layers_to_capture=LAYERS, hidden_size=H), _drafter() + ) + assert out is hs + + +@pytest.mark.parametrize("normed", [set(), {1, 29}, {29, 56}, {1, 56}]) +def test_raw_layers_get_their_own_fc_norm(normed) -> None: + hs = torch.randn(4, H * len(LAYERS)) + ref = hs.clone() + out = _worker(normed)._m3_renorm_raw_captures( + hs, SimpleNamespace(layers_to_capture=LAYERS, hidden_size=H), _drafter() + ) + assert torch.equal(hs, ref) # the capture buffer itself is not modified + for i, layer_id in enumerate(LAYERS): + block, src = out[:, i * H : (i + 1) * H], ref[:, i * H : (i + 1) * H] + assert torch.equal(block, src if layer_id in normed else src * (i + 2)), layer_id + + +def test_arm_resets_the_record_and_the_fork_result() -> None: + w = _worker(set(LAYERS)) + w._m3_fc_hidden = object() + w._m3_fc_early_stream = None # no fork pending + w._m3_join_pending = lambda: None # (tp8) nothing else pending + w.arm_capture_norm( + None, None, None, None + ) # no drafter: resets the per-forward state and returns + assert w._m3_capture_normed == set() + assert w._m3_fc_hidden is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py new file mode 100644 index 000000000000..3299a8432362 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The worker's capture target (the capture norm in the boundary all-reduce): the capture buffer's column block of a +captured layer with its fc_norm weight and eps; None without an armed capture norm, for a layer that is not captured, +past the buffer's rows, or when the loaded all-reduce op has no capture epilogue (then the capture_norm kernel runs).""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_mc_allreduce, +) + +H = 64 +LAYERS = (2, 30, 58) + + +def _setup(armed=True): + hidden = torch.arange(8 * 3 * H, dtype=torch.float32).reshape(8, 3 * H).to(torch.bfloat16) + weights = torch.randn(3, H).to(torch.bfloat16) + meta = SimpleNamespace(hidden_states=hidden, hidden_size=H, layers_to_capture=LAYERS) + worker = SimpleNamespace(_m3_capture_norm=(weights, 1e-5) if armed else None) + return worker, meta, hidden, weights + + +def _target(worker, meta, layer_id, num_tokens): + return modeling.M3Eagle3Worker.capture_target(worker, meta, layer_id, num_tokens) + + +@pytest.fixture +def supported(monkeypatch): + monkeypatch.setattr(m3_mc_allreduce, "capture_supported", lambda: True) + + +def test_column_block_of_each_captured_layer(supported) -> None: + worker, meta, hidden, weights = _setup() + for i, layer_id in enumerate(LAYERS): + out, w, eps = _target(worker, meta, layer_id, 4) + assert out.shape == (4, H) + assert out.data_ptr() == hidden[:4, i * H : (i + 1) * H].data_ptr() + assert out.stride() == hidden.stride() + assert torch.equal(w, weights[i]) + assert eps == 1e-5 + + +def test_none_when_not_capturable(supported) -> None: + worker, meta, _, _ = _setup() + assert _target(worker, meta, 3, 4) is None # not a captured layer + assert _target(worker, meta, LAYERS[0], 9) is None # past the buffer's rows + unarmed, meta2, _, _ = _setup(armed=False) + assert _target(unarmed, meta2, LAYERS[0], 4) is None + + +def test_none_without_the_capture_epilogue(monkeypatch) -> None: + monkeypatch.setattr(m3_mc_allreduce, "capture_supported", lambda: False) + worker, meta, _, _ = _setup() + assert _target(worker, meta, LAYERS[0], 4) is None + + +def test_capture_supported_reads_the_op_schema() -> None: + # The schema of the op the boundaries call: this target's copy (m3_mc_allreduce.OWN_OP), else the image's. + name = ( + "minimax_m3_sm107_tp8_mc_allreduce_norm" + if m3_mc_allreduce.own_op() + else "minimax_m3_sm107_tp2_mc_allreduce_norm" + ) + op = getattr(torch.ops.trtllm, name, None) + want = op is not None and any(a.name == "cap_out" for a in op.default._schema.arguments) + assert m3_mc_allreduce.capture_supported() is want + + +def test_own_op_follows_the_kernel_constant(monkeypatch) -> None: + monkeypatch.setattr(m3_mc_allreduce, "OWN_OP", False) + assert m3_mc_allreduce.own_op() is False + monkeypatch.setattr(m3_mc_allreduce, "OWN_OP", True) + loaded = getattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_mc_allreduce_norm", None) is not None + assert m3_mc_allreduce.own_op() is loaded diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py new file mode 100644 index 000000000000..9a82b7505885 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The lazily created exchange workspaces refuse to be created inside a CUDA-graph capture (K3 audit #150, check #10). + +The draft-step and sampled-accept workspaces are created collectively (both ranks, a multicast buffer) on the first +eager call. Creating one inside a capture would allocate and synchronize mid-capture, so a capture that finds it missing +reports a miss (``kernels/m3_report.capture_missed``) and gets None (the caller then takes its stock path) instead. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( + m3_draft_step, + m3_report, + m3_sampled_accept, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU" +) + + +def _refuse(*args, **kwargs): + raise AssertionError("a workspace must not be created inside a CUDA-graph capture") + + +@pytest.mark.parametrize( + "module, creator", + [ + (m3_draft_step, "create_workspace"), + (m3_sampled_accept, "SampledAcceptWorkspace.create"), + ], + ids=["draft_step", "sampled_accept"], +) +def test_a_capture_without_the_workspace_misses_instead_of_creating_it( + monkeypatch, module, creator +): + monkeypatch.setattr(module, "_WS", None) + owner, _, attr = creator.rpartition(".") + monkeypatch.setattr(getattr(module, owner) if owner else module, attr, _refuse) + misses = [] + monkeypatch.setattr(m3_report, "capture_missed", lambda *args: misses.append(args)) + + device = torch.device("cuda", torch.cuda.current_device()) + found = [] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + found.append(module.workspace(None, device)) + torch.cuda.current_stream().wait_stream(stream) + + assert found == [None] + assert module._WS is None + assert len(misses) == 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..8e8982524c8b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_chunk_norm entry (the Eagle3 fc_norm over the captured layers).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_chunk_norm import ( # noqa: E501 + m3_chunk_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, c, h = spec["m"], spec["c"], spec["h"] + buf = randn((m, c * h + spec["extra_cols"]), g, 2.0) + src = buf[:, : c * h + spec["extra_cols"]] + w = (1.0 + randn((c, h), g, 0.1).float()).to(torch.bfloat16) + dst = torch.full((m, c * h), 7.0, dtype=torch.bfloat16, device="cuda") + return (src, w, EPS, dst, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + + +def test_guard_refuses_a_strided_weight() -> None: + src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() + wt.copy_(w) + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(src, wt, eps, dst, pdl) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py new file mode 100644 index 000000000000..3bcee7c3345a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the fc fork's counter wait (modeling's DI_FC_FLAG, kernels/m3_fc_wait): the one-thread wait kernel +queued on the main stream releases its successor only after the fork's bump on another stream (so the successor reads +what the fork wrote, with no stream join before it), eagerly and on CUDA-graph replays; it counts an overrun instead +of hanging when no bump comes.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_fc_wait, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + +N = 1 << 16 + + +def _wait(counts, spins=m3_fc_wait.MAX_SPINS): + m3_fc_wait.fc_wait_kernel[(1,)](counts, MAX_SPINS=spins, num_warps=1, launch_pdl=True) + + +def _step(counts, aux, src, dst, value, delay_cycles): + """One fork / wait / consume: the fork (after a delay) writes ``value`` into src and bumps the counter; the main + stream waits on the counter, copies src to dst, then joins the fork (deferred, as the worker does).""" + main = torch.cuda.current_stream() + aux.wait_stream(main) + with torch.cuda.stream(aux): + torch.cuda._sleep(delay_cycles) + src.fill_(value) + counts[:1].add_(1) + _wait(counts) + dst.copy_(src) + main.wait_stream(aux) + + +@pytest.mark.parametrize("delay_cycles", [0, 200_000, 2_000_000]) +def test_wait_releases_after_the_fork(delay_cycles) -> None: + counts = torch.zeros(3, dtype=torch.int64, device="cuda") + src = torch.zeros(N, dtype=torch.float32, device="cuda") + dst = torch.empty_like(src) + aux = torch.cuda.Stream() + for it in range(1, 9): + _step(counts, aux, src, dst, float(it), delay_cycles) + torch.cuda.synchronize() + assert bool((dst == float(it)).all()), f"iteration {it}: read before the fork finished" + assert counts.tolist() == [8, 8, 0] + + +def test_wait_releases_on_graph_replays() -> None: + counts = torch.zeros(3, dtype=torch.int64, device="cuda") + src = torch.zeros(N, dtype=torch.float32, device="cuda") + dst = torch.empty_like(src) + val = torch.zeros(1, dtype=torch.float32, device="cuda") + aux = torch.cuda.Stream() + cap = torch.cuda.Stream() + + def body(): + main = torch.cuda.current_stream() + aux.wait_stream(main) + with torch.cuda.stream(aux): + torch.cuda._sleep(1_000_000) + src.copy_(val.expand(N)) + counts[:1].add_(1) + _wait(counts) + dst.copy_(src) + main.wait_stream(aux) + + with torch.cuda.stream(cap): + body() # eager warm-up (compiles the wait kernel outside capture) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(cap): + with torch.cuda.graph(graph, stream=cap): + body() + torch.cuda.synchronize() + base = counts.tolist() + for it in range(1, 11): + val.fill_(float(it)) + graph.replay() + torch.cuda.synchronize() + assert bool((dst == float(it)).all()), f"replay {it}: read before the fork finished" + assert counts.tolist() == [base[0] + 10, base[1] + 10, 0] + + +def test_wait_without_a_fork_counts_an_overrun() -> None: + counts = torch.zeros(3, dtype=torch.int64, device="cuda") + _wait(counts, spins=1000) + torch.cuda.synchronize() + assert counts.tolist() == [0, 1, 1] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..e1b7c9c625bb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_input entry (embedding + both RMSNorms + concat in one launch), against its reference and +against the in-tree drafter's modules.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_input import ( # noqa: E501 + m3_draft_input as op, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + VOCAB, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, h = spec["m"], spec["h"] + table = randn((VOCAB, h), g) + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + if spec["invalid_id"]: + ids[0] = -1 + hidden = randn((4 * m if spec["gather"] else m, h), g) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 2 if spec["gather"] else None + w_e = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + w_h = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + out = torch.full((m, 2 * h), 7.0, dtype=torch.bfloat16, device="cuda") + residual = torch.full((m, h), 7.0, dtype=torch.bfloat16, device="cuda") + return (ids, table, hidden, rows, w_e, w_h, EPS, EPS, out, residual, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Normalized halves within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic in a + different summation order, rsqrtf); the gathered residual and the zero embedding of an invalid id exact.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out, residual = op(*args) + ref_out, ref_res = op.reference(*args) + torch.cuda.synchronize() + assert_bf16_close(out, ref_out) + assert torch.equal(residual, ref_res) + if cell.spec["invalid_id"]: + assert not bool(out[0, : cell.spec["h"]].float().any()) + + +def test_guard_refuses_rows_past_hidden() -> None: + gather = next(c for c in op.CELLS if c.spec["gather"]) + args = list(_case(gather.spec, 1)) + args[3] = args[3] + 100 + with pytest.raises(AssertionError, match="rows must index hidden"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("m,gather", [(4, False), (32, False), (8, True)], ids=str) +def test_the_in_tree_draft_input(m, gather) -> None: + """Against the in-tree drafter's own modules: the Embedding lookup, then the layer's input_layernorm and + hidden_norm (in-tree RMSNorm) and the concat of Eagle3DecoderLayer.forward. Within one bf16 ulp on at most 0.1% of + the elements (the RMS statistic summed in another order); the residual (the carried hidden rows) exact.""" + g = gen(seed_of(f"in-tree draft input {m} {gather}")) + h = HIDDEN + embed = Embedding(VOCAB, h, dtype=torch.bfloat16).cuda() + embed.weight.data = randn((VOCAB, h), g) + norms = [] + for _ in range(2): + norm = RMSNorm(hidden_size=h, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + norms.append(norm) + input_layernorm, hidden_norm = norms + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + hidden = randn((4 * m if gather else m, h), g, 2.0) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 3 if gather else None + carried = hidden[rows] if gather else hidden + stock = torch.cat([input_layernorm(embed(ids)), hidden_norm(carried)], dim=-1) + out, residual = m3_spec_tail.draft_input( + ids, embed.weight, hidden, rows, input_layernorm.weight, EPS, hidden_norm.weight, EPS + ) + torch.cuda.synchronize() + assert_bf16_close(out, stock) + assert torch.equal(residual, carried) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_push_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_push_take.py new file mode 100644 index 000000000000..8f0896d6fca6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_push_take.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter boundary push bookkeeping (``M3Eagle3Draft.arm_push`` / ``take_pushed``, ``DRAFT_AR_PUSH``): the +boundary all-reduce runs prepushed only when the armed call pushed, and only on that call's own output; any other +input raises instead of reducing rows the exchange does not hold. Host state only (no GPU).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) + + +@pytest.fixture +def draft(): + """A drafter with only its push bookkeeping (the module and its weights are not built).""" + d = object.__new__(modeling.M3Eagle3Draft) + d._push_arm, d._pushed = None, 0 + return d + + +def test_not_pushed_is_false_and_disarms(draft) -> None: + draft._push_arm = (1, 2, 3, 0, 8, 8) + assert draft.take_pushed(torch.zeros(4)) is False + assert draft._push_arm is None + + +def test_pushed_output_is_true_once(draft) -> None: + out = torch.zeros(4, 8) + draft._pushed = out.data_ptr() + assert draft.take_pushed(out) is True + assert draft._pushed == 0 + assert draft.take_pushed(out) is False + + +def test_other_input_raises(draft) -> None: + out = torch.zeros(4, 8) + other = out.clone() + draft._pushed = out.data_ptr() + with pytest.raises(RuntimeError, match="not the boundary all-reduce's input"): + draft.take_pushed(other) + assert draft._pushed == 0 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py new file mode 100644 index 000000000000..5ae74445fe33 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_silu entry (the drafter MLP's plain SwiGLU as one C++ launch): bit for bit the Triton +silu_and_mul at its defaults, over every bf16 gate and every bf16 up pattern.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_silu import ( # noqa: E501 + m3_draft_silu as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" + +ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) + + +def _rows(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """[rows, 2d] gate | up from equally long flat gate and up values, d = 2304 (the TP8 drafter shard).""" + d = 2304 + n = gate.numel() + rows = -(-n // d) + pad = rows * d - n + g = torch.cat([gate, gate[:pad]]).view(rows, d) + u = torch.cat([up, up[:pad]]).view(rows, d) + return torch.cat([g, u], dim=1).contiguous() + + +def _assert_bitwise(x: torch.Tensor, what: str) -> None: + got = op(x) + assert got is not None, what + ref = op.reference(x) + torch.cuda.synchronize() + diff = got.view(torch.int16) != ref.view(torch.int16) + assert not bool(diff.any()), f"{what}: {int(diff.sum())} of {diff.numel()} differ" + + +@pytest.mark.parametrize("up", [0.0, 1.0, -1.0, 3.5, -0.0, 1e-3], ids=lambda v: f"up={v}") +def test_every_gate_pattern_is_bitwise_the_triton_kernel(up) -> None: + gate = ALL.cuda() + _assert_bitwise(_rows(gate, torch.full_like(gate, up)), f"gate sweep, up {up}") + + +@pytest.mark.parametrize("gate", [1.0, -1.0, 0.5, -3.0, 20.0, -0.0], ids=lambda v: f"gate={v}") +def test_every_up_pattern_is_bitwise_the_triton_kernel(gate) -> None: + up = ALL.cuda() + _assert_bitwise(_rows(torch.full_like(up, gate), up), f"up sweep, gate {gate}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed) -> None: + g = gen(100 + seed) + n = 1 << 22 + gate = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + up = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + _assert_bitwise(_rows(gate.view(torch.bfloat16), up.view(torch.bfloat16)), f"seed {seed}") + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The drafter's gate_up output range (and past it), in a CUDA graph as the drafter runs it.""" + g = gen(seed_of(cell.why)) + rows, d = cell.spec["rows"], cell.spec["d"] + x = (4.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :6] = torch.tensor([60.0, -60.0, 0.0, -0.0, 88.0, -88.0], dtype=torch.bfloat16) + certifying(op) + ref = op.reference(x) + eager = op(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(x) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(eager.view(torch.int16), ref.view(torch.int16)) + assert torch.equal(captured.view(torch.int16), ref.view(torch.int16)) + + +def test_declines_what_it_does_not_cover() -> None: + x = torch.zeros(2, 2 * 2304, dtype=torch.bfloat16, device="cuda") + assert op(x[:, 8:]) is None # strided rows + assert ( + op(torch.zeros(2, 2 * 12, dtype=torch.bfloat16, device="cuda")) is None + ) # d not a multiple of 8 + assert op(x.float()) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..58e3c0a89b6d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,747 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_step entry (one fused greedy draft pass: the BF16 draft head, streamed from its lossless +12-bit coding or as BF16 fragments, the cross-rank argmax and the draft loop's bookkeeping), and its tokens against the +in-tree drafter's head and greedy draft sampling. + +The in-tree path: each rank's LM head shard GEMM (BF16 out), ``LogitsProcessor``'s float cast, then either the argmax +of the gathered logits (``torch.argmax``) or ``greedy_sample_draft_with_tp_gather``'s per-rank ``torch.max`` and the +pick among the ranks' (index, value) pairs. The edge cases use integer data, whose fp32 sums are exact in any order, +so the kernel and the in-tree path see the same BF16 logits and must agree bit for bit; every in-tree comparison runs +with the coded head (the default) and with the BF16 fragments, and on integer data the coded head's winning keys are +the BF16 fragments' bit for bit. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.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_107__tp8.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_LEN, + DRAFT_SHARD, + HIDDEN, + MAX_TOKENS, + bf16_fragments, + from_fragments, + head_weight, + max_tile_pieces, + shard_logits, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_step import ( # noqa: E501 + m3_draft_step as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + WORLD, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" + +T = DRAFT_LEN + 1 +_LOW32 = 0xFFFFFFFF + +# ------------------------------------------------------------------ the per-row rule against the in-tree token +# +# The kernel and the in-tree head (a cuBLAS BF16 GEMM) both round an FP32 sum of the same exact BF16 products to BF16, +# in different summation orders. A row's token can differ from the in-tree one only where some candidate's FP32 sum +# can land on either side of a BF16 rounding midpoint. Each row is held to: +# +# kernel token == in-tree token, or the row is provably undecidable: with the fp64 reference logit L_i (exact BF16 +# products, fp64 sum) and the FP32-order bound E_i = gamma_K * sum_k |w_ik * x_k|, the BF16 value of any FP32 sum of +# candidate i lies in [lo_i, hi_i] = [bf16(L_i - E_i), bf16(L_i + E_i)]; the kernel's token must win for some +# choice of values within those intervals (larger value, then the lower global id), and a BF16 rounding midpoint +# must lie within E of the kernel's or the in-tree token's L. +# +# gamma_K = K * 2^-22 is la-xh2's bound (4x the round-to-nearest unit roundoff 2^-24 per addition: any summation +# order, and a truncating tensor-core accumulator); K counts the row's non-zero products (adding an exact zero does +# not round). lo / hi take one more FP32 step outward before the BF16 rounding, which covers rounding L -+ E to FP32. +# Each undecidable row is reported as a warning with its midpoint distances, and each test reports its row count; any +# other mismatch fails. +GAMMA_UNIT = 2.0**-22 + + +def _bf16_outward(v: torch.Tensor, direction: float) -> torch.Tensor: + f = v.float() + f = torch.nextafter(f, torch.full_like(f, direction)) + b = f.to(torch.bfloat16).float() + return torch.where(b == 0, torch.zeros_like(b), b) + + +def bounds(h: torch.Tensor, w: torch.Tensor): + """(fp64 logits L, bound E, lo, hi) [rows, N] of the head input rows h against the head rows w.""" + cols = (h != 0).any(dim=0) # zero columns add exact zeros + hd, wd = h[:, cols].double(), w[:, cols].double() + logits = hd @ wd.t() + bound = (h != 0).sum(dim=1, keepdim=True).double() * GAMMA_UNIT * (hd.abs() @ wd.abs().t()) + del wd + inf = float("inf") + return logits, bound, _bf16_outward(logits - bound, -inf), _bf16_outward(logits + bound, inf) + + +def midpoint_distance(v: float) -> float: + """Distance of v to the nearest BF16 rounding midpoint.""" + b = torch.tensor(v, dtype=torch.float64).float().to(torch.bfloat16).float() + if float(b) == 0.0: + return abs(v) + ulp = 2.0 ** (int(torch.frexp(b).exponent) - 8) + return min(abs(v - (float(b) + ulp / 2)), abs(v - (float(b) - ulp / 2))) + + +def _winnable(tok: int, ids: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Whether candidate ``tok`` wins the row for some values within the candidates' [lo, hi].""" + k = int((ids == tok).nonzero()[0, 0]) + others = ids != tok + beat = (hi[k] > lo[others]) | ((hi[k] == lo[others]) & (tok < ids[others])) + return bool(beat.all()) + + +def check_row(name: str, kernel_tok: int, stock_tok: int, ids, logits, bound, lo, hi) -> bool: + """The per-row rule for one row over the candidates ``ids`` [C] (logits / bound / lo / hi [C]); True when the + row is undecidable (reported as a warning), False when the tokens match; a mismatch outside the bound fails.""" + if kernel_tok == stock_tok: + return False + assert _winnable(kernel_tok, ids, lo, hi), ( + f"{name}: token {kernel_tok} differs from the in-tree {stock_tok} outside the FP32-order bound" + ) + ik = int((ids == kernel_tok).nonzero()[0, 0]) + i_stock = int((ids == stock_tok).nonzero()[0, 0]) + lk, ls, ek, es = ( + float(logits[ik]), + float(logits[i_stock]), + float(bound[ik]), + float(bound[i_stock]), + ) + dk, ds = midpoint_distance(lk), midpoint_distance(ls) + assert dk <= ek or ds <= es, f"{name}: neither logit is within its bound of a BF16 midpoint" + warnings.warn( + f"[draft-step undecidable] {name}: kernel {kernel_tok} (L {lk:.9g}, midpoint distance {dk:.3e}, bound " + f"{ek:.3e}) vs in-tree {stock_tok} (L {ls:.9g}, midpoint distance {ds:.3e}, bound {es:.3e})" + ) + return True + + +def report(name: str, rows: int, undecidable: int) -> None: + warnings.warn(f"[draft-step rows] {name}: {rows} rows, {undecidable} undecidable") + + +def key_value(keys: torch.Tensor) -> torch.Tensor: + """The fp32 value a key orders by (the inverse of the kernel's order-preserving image).""" + u = (keys >> 32) & _LOW32 + bits = torch.where(u >= 0x80000000, u - 0x80000000, _LOW32 - u) + bits = torch.where(bits >= 2**31, bits - 2**32, bits) # as a signed 32-bit pattern + return bits.to(torch.int32).view(torch.float32) + + +def key_index(keys: torch.Tensor) -> torch.Tensor: + return _LOW32 - (keys & _LOW32) + + +def _bookkeeping(nreq: int, step: int, n_x_rows: int, g: torch.Generator): + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if step == 0 + else None + ) + seq_lens = torch.full((nreq,), T, dtype=torch.int32, device="cuda") if step == 0 else None + next_new = torch.full((nreq, T), -5, dtype=torch.int32, device="cuda") + pos_src = torch.randint(100, 3000, (n_x_rows,), generator=g, device="cuda", dtype=torch.int32) + kv_lens = torch.randint(200, 3000, (nreq,), generator=g, device="cuda", dtype=torch.int32) + ids_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + pos_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + return num_acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next + + +def _head(w: torch.Tensor, coded: bool): + """(wf, (exc_off, exc_pos, exc_w, exp_base)) of a BF16 shard: its 12-bit coding (checked to round-trip) or its + BF16 fragments.""" + if not coded: + return bf16_fragments(w), (None, None, None, 0) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + return packed, (exc_off, exc_pos, exc_w.view(-1), base) + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=(None, None, None, 0), +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + coded head's (exc_off, exc_pos, exc_w, exp_base).""" + acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next = _bookkeeping( + nreq, step, x.shape[0], g + ) + if num_acc is None: + num_acc = acc + if uc is None: + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + out_keys = torch.zeros(nreq, dtype=torch.int64, device="cuda") + tiles = wf.shape[0] if coded[0] is None else coded[0].numel() - 1 + return (x, x_rows, nreq, wf, rank * tiles * 16, keys, counters, uc, uc.data_ptr(), rank, world, + next_new[:, 1:], step, ids_next, pos_src, pos_next, kv_lens, num_acc, seq_lens, DRAFT_LEN, num_sms(), + True, out_keys, *coded) # fmt: skip + + +def _fixed(args) -> bool: + """fixed_order for these op arguments as dispatch picks it (kernels/m3_draft_step._fixed_order): the warp-order + sum where the launch grid lets a tile take more than 2 warp pieces (D-FIX; the TP8 shard on 212 SMs does).""" + x, wf, num_ctas, exc_off = args[0], args[3], args[20], args[23] + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + return max_tile_pieces(tiles, num_ctas, x.shape[1] // (16 if exc_off is None else 64)) > 2 + + +def _case(spec, seed): + g = gen(seed) + nreq, step, rank, world, vs, k = ( + spec["nreq"], + spec["step"], + spec["rank"], + spec["world"], + spec["vs"], + spec["k"], + ) + pass0 = step == 0 + x = randn((nreq * T if pass0 else nreq, k), g) + wf, coded = _head(randn((vs, k), g, 0.02), spec["coded"]) + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if pass0 + else None + ) + x_rows = (torch.arange(nreq, device="cuda") * T + num_acc.long() - 1) if pass0 else None + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + h = x[x_rows] if pass0 else x + own = shard_logits(h, wf, *coded).max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=MAX_TOKENS, ties=spec.get("ties", False) + ) + return _args(x, x_rows, nreq, wf, rank, world, step, g, uc, keys, counters, num_acc, coded) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The bookkeeping is exact on every row (next positions, KV and sequence lengths). Each row's token is the + in-tree one (this rank's BF16 shard logits from the head's GEMM, their first maximum, against the peers' keys) or + the row is provably undecidable (check_row); the winning key's index is the token, and its value lies in the + token's [lo, hi] (or is the peer's pushed key when a peer wins).""" + args = _case(cell.spec, seed_of(cell.why)) + x, x_rows, nreq, wf, offset, _, counters, uc = args[:8] + rank, world = args[9], args[10] + half = int(counters[1].item()) & 1 + pushed = ( + uc[half * MAX_TOKENS * world : (half * MAX_TOKENS + nreq) * world].view(nreq, world).clone() + ) + h = x[x_rows] if x_rows is not None else x[:nreq] + w = head_weight(wf, x.shape[1], *args[23:27]) + best, idx = local_argmax(torch.nn.functional.linear(h, w).float(), offset) + stock = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, MAX_TOKENS) + certifying(op) + ref = op.reference(*args) + op(*args, fixed_order=_fixed(args)) + torch.cuda.synchronize() + assert args[6].tolist() == [0, 1] and not bool(args[5].any()), "counters / keys not re-armed" + step = args[12] + others = [c for c in range(args[11].shape[1]) if c != step] + assert torch.equal(args[11][:, others], ref["draft_out"][:, others]) + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + assert torch.equal(args[15], ref["pos_next"]) + assert torch.equal(args[16], ref["kv_lens"]) + if args[18] is not None: + assert torch.equal(args[18], ref["seq_lens"]) + # Candidates per row: this rank's head rows (fp64 bounds) and each peer's pushed key (an exact value). + logits, bound, lo, hi = bounds(h, w) + own_ids = offset + torch.arange(w.shape[0], device="cuda") + peers = [p for p in range(world) if p != rank] + undecidable = 0 + for r in range(nreq): + pv = key_value(pushed[r, peers]).double() + ids = torch.cat([own_ids, key_index(pushed[r, peers])]) + undecidable += check_row( + f"{cell.why[:40]} row {r}", + int(args[13][r]), + int(stock[r]), + ids, + torch.cat([logits[r], pv]), + torch.cat([bound[r], torch.zeros_like(pv)]), + torch.cat([lo[r], pv.float()]), + torch.cat([hi[r], pv.float()]), + ) + report(cell.why[:40], nreq, undecidable) + out_keys = args[22] + assert torch.equal(key_index(out_keys), args[13].long()), "the winning key's index is the token" + got = key_value(out_keys) + for r in range(nreq): + c = int(args[13][r]) - offset + if 0 <= c < w.shape[0]: + assert float(lo[r, c]) <= float(got[r]) <= float(hi[r, c]), ( + f"row {r}: winning value outside [lo, hi]" + ) + else: + assert int(out_keys[r]) in pushed[r].tolist(), ( + f"row {r}: a peer won with a key it did not push" + ) + + +def test_fragments_round_trip() -> None: + g = gen(seed_of("fragments")) + w = randn((64, 96), g) + wf = bf16_fragments(w) + assert wf.shape == (4, 6, 256) + assert torch.equal(from_fragments(wf), w) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 3) # K 96: no coding, the BF16 fragments + assert not prepared.coded + assert torch.equal(prepared.wf, wf) and prepared.vocab_offset == 3 * 64 and prepared.weight is w + + +# ------------------------------------------------------------------ against the in-tree drafter, bit for bit + + +def _stock_tokens(shards: list, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The in-tree greedy draft tokens of the rows x over the vocab-sharded BF16 head: (argmax of the gathered float + logits, greedy_sample_draft_with_tp_gather's per-rank (index, value) pick).""" + logits = [torch.nn.functional.linear(x, w).float() for w in shards] + gathered = torch.argmax(torch.cat(logits, dim=-1), dim=-1).to(torch.int32) + combined = [] + for rank, shard_logit in enumerate(logits): + worker = SimpleNamespace(mapping=SimpleNamespace(tp_rank=rank)) + combined.append(Eagle3OneModelWorker._get_local_max_and_combined(worker, shard_logit)) + pick = Eagle3OneModelWorker._get_draft_tokens_from_gathered + pick = getattr(pick, "_torchdynamo_orig_callable", pick) + tp_gather = pick(None, torch.cat(combined, dim=-1)).to(torch.int32) + return gathered, tp_gather + + +def _fused_tokens(shards: list, x: torch.Tensor, coded: bool, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's four-rank pass with + its peers' keys in its slots, as their multicast pushes would put them.""" + g = gen(seed_of("fused tokens")) + nreq = x.shape[0] + heads = [_head(w, coded) for w in shards] + own_keys = [] + for rank, (wf, cd) in enumerate(heads): + args = list(_args(x, None, nreq, wf, rank, 1, step, g, coded=cd)) + args[9] = 0 # one-rank run: this rank's slot only, global ids of `rank` + op(*args, fixed_order=_fixed(args)) + own_keys.append(args[22].clone()) + tokens = [] + for rank, (wf, cd) in enumerate(heads): + uc, keys, counters = head_workspace(WORLD, slot_rows=MAX_TOKENS) + for peer in range(WORLD): + if peer != rank: + for t in range(nreq): + uc[t * WORLD + peer] = own_keys[peer][t] + args = _args(x, None, nreq, wf, rank, WORLD, step, g, uc, keys, counters, coded=cd) + with validating(op): + op(*args, fixed_order=_fixed(args)) + torch.cuda.synchronize() + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + tokens.append(args[13].clone()) + return tokens + + +def _integer_head(g: torch.Generator, rows: int = WORLD * DRAFT_SHARD) -> torch.Tensor: + """A [rows, 6144] head of small integers over the first 4 columns (zero elsewhere): its logits against integer + rows are exact integers in any summation order, many of which share a BF16 value.""" + w = torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, :4] = torch.randint(-60, 61, (rows, 4), generator=g, device="cuda").to(torch.bfloat16) + return w + + +def _shards(w: torch.Tensor) -> list: + return [w[r * DRAFT_SHARD : (r + 1) * DRAFT_SHARD].contiguous() for r in range(WORLD)] + + +def _assert_stock(w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None) -> None: + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + if expect is not None: + assert gathered.tolist() == expect, (gathered.tolist(), expect) + for rank, tokens in enumerate(_fused_tokens(shards, x, coded)): + assert torch.equal(tokens, gathered), ( + f"rank {rank}: {tokens.tolist()} vs stock {gathered.tolist()}" + ) + + +def _ones_rows(nreq: int) -> torch.Tensor: + x = torch.zeros((nreq, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[:, :4] = 1.0 + return x + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_forced_bf16_ties_take_the_lowest_global_id(coded) -> None: + """A head row copied onto every rank that wins every row of ones (logit 244, every other at most 240): the lowest + global id must win, inside a shard and across ranks. Then integer rows in [-2, 2], whose logits (up to 480 in + magnitude, BF16 spacing 2 above 256) collide often, exactly or after BF16 rounding.""" + g = gen(seed_of("forced ties")) + w = _integer_head(g) + win = torch.tensor([61.0, 61.0, 61.0, 61.0], device="cuda").to(torch.bfloat16) + for gid in (DRAFT_SHARD + 5000, DRAFT_SHARD + 77, DRAFT_SHARD + 9000, 23750): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), coded, expect=[23750] * 5) + # Random integer rows: ties after BF16 rounding happen naturally. + x = torch.randint(-2, 3, (8, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + x[:, 4:] = 0 + _assert_stock(w, x, coded) + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_signed_zero_maxima_take_the_lowest_global_id(coded) -> None: + """Rows whose maximum is zero: an all-zero head input (every logit +-0) and logits <= 0 with zeros by cancellation + on several ranks. -0 and +0 compare equal, so the lowest global id reaching zero wins.""" + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 0] = -1.0 + w[:, 1] = -1.0 + for gid in (DRAFT_SHARD + 5, DRAFT_SHARD + 40000, DRAFT_SHARD + 70007, 24999): + w[gid, 0] = 1.0 + w[gid, 1] = -1.0 + x = torch.zeros((3, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[1, :2] = 1.0 + x[2, :2] = 2.0 + x[2, 2] = -1.0 # w[:, 2] is zero: a -0 product on every row + _assert_stock(w, x, coded, expect=[0, 24999, 24999]) + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + (1000, 1001, "lower"), # 1001 rounds to 1000 in BF16: a tie, the lower id wins + (1001, 1002, "lower"), # 1001 and the tie-to-even midpoint 1002 both round to 1000 + (1000, 1004, "higher"), # adjacent BF16 values: the larger wins + (1003, 1000, "lower"), # 1003 rounds up to 1004 > 1000 + ], + ids=["one_fp32_step", "midpoint", "adjacent_bf16", "rounds_up"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where, coded) -> None: + """Top-2 logits one to four fp32 steps apart: the winner is the one the BF16 logits decide (the in-tree logits' + dtype), not the fp32 sum. The pair sits inside rank 1's shard or on ranks 0 and 1; every other logit is at most + 240 in magnitude.""" + g = gen(seed_of(f"near top2 {lower} {higher} {where}")) + w = _integer_head(g) + if where == "in_shard": + lo_id, hi_id = DRAFT_SHARD + 100, DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD - 1000, DRAFT_SHARD + 3 + for gid, value in ((lo_id, lower), (hi_id, higher)): + q, r = divmod(value, 4) # four integers <= 251, exact in BF16 + w[gid, :4] = torch.tensor([q + (i < r) for i in range(4)], device="cuda").to(torch.bfloat16) + expect = lo_id if winner == "lower" else hi_id + _assert_stock(w, _ones_rows(2), coded, expect=[expect, expect]) + + +# ------------------------------------------------------------------ non-integer rows at a BF16 rounding midpoint + +#: A BF16 rounding midpoint (between 6.0 and 6.03125; a tie there rounds to the even 6.0) and the FP32 spacing there. +MIDPOINT = 6.0 + 2.0**-6 +FP32_ULP = 2.0**-21 + + +def _near_midpoint_head(g: torch.Generator, targets: dict): + """A full-vocabulary head [8 * 25008, 6144] and one head-input row x with 65 non-zero entries: 1, 2^-12, 2^-20 + and 2^-28 on columns 0..3 and 61 random non-integer BF16 values. Every head row carries random non-integer BF16 + values on those 61 columns (logits about 0.2 in magnitude); each row of ``targets`` (global id -> fp64 value) also + gets columns 0..3 set so its fp64 logit is the target to within 2^-30 (a coarse value, then three finer + corrections, each BF16-rounded).""" + x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device="cuda") + scales = (1.0, 2.0**-12, 2.0**-20, 2.0**-28) + for col, scale in enumerate(scales): + x[0, col] = scale + x[0, 4:65] = randn((61,), g, 0.5) + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 4:65] = randn((WORLD * DRAFT_SHARD, 61), g, 0.05) + xd = x[0].double() + for gid, target in targets.items(): + rest = float(target - (w[gid].double() * xd).sum()) + for col, scale in enumerate(scales): + v = torch.tensor(rest / scale, dtype=torch.float64).to(torch.bfloat16) + w[gid, col] = v + rest -= float(v) * scale + assert abs(rest) <= 2.0**-30, (gid, rest) + return w, x + + +def _check_fused_against_stock( + name: str, w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None +) -> int: + """Each rank's fused tokens against the in-tree tokens under the per-row rule (full-vocabulary candidates); + ``expect`` (a token per row) additionally pins rows the bound decides. Returns the undecidable rows.""" + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + fused = _fused_tokens(shards, x, coded) + for tokens in fused[1:]: + assert torch.equal(tokens, fused[0]), "the ranks disagree on the global token" + logits, bound, lo, hi = bounds(x, w) + ids = torch.arange(w.shape[0], device="cuda") + undecidable = 0 + for r in range(x.shape[0]): + undecidable += check_row( + f"{name} row {r}", + int(fused[0][r]), + int(gathered[r]), + ids, + logits[r], + bound[r], + lo[r], + hi[r], + ) + if expect is not None: + assert int(gathered[r]) == expect[r] and int(fused[0][r]) == expect[r], ( + name, + int(gathered[r]), + int(fused[0][r]), + expect[r], + ) + others = ids != expect[r] + assert _winnable(expect[r], ids, lo[r], hi[r]), ( + f"{name}: the bound does not let {expect[r]} win" + ) + rivals = [int(i) for i in ids[others][hi[r][others] >= lo[r][expect[r]]]] + assert not any(_winnable(i, ids, lo[r], hi[r]) for i in rivals), ( + f"{name}: the row is not decided" + ) + report(name, x.shape[0], undecidable) + return undecidable + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + ( + -2, + 2, + None, + ), # straddling the midpoint by 2 FP32 ulps: either side is in-tree up to the order + (2, -2, None), + (0, -2, None), # exactly on the midpoint (ties to the even 6.0) against just below it + ( + -2048, + 2048, + "higher", + ), # ~1e-3 either side, well outside the bound: decided by BF16 rounding + ( + -4096, + -2048, + "lower", + ), # both below the midpoint, ~1e-3 apart: a BF16 tie the lower id wins + ], + ids=["straddle", "straddle_reversed", "on_midpoint", "decided_up", "decided_tie"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where, coded) -> None: + """Non-integer rows whose top-2 fp64 logits sit the given number of FP32 ulps from a BF16 rounding midpoint, inside + rank 1's shard or on ranks 0 and 1, every other logit far below. The per-row rule holds on every row; outside the + bound the token is the one the BF16 rounding decides, bit for bit.""" + g = gen(seed_of(f"midpoint {lower} {higher} {where}")) + if where == "in_shard": + lo_id, hi_id = DRAFT_SHARD + 100, DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD - 1000, DRAFT_SHARD + 3 + w, x = _near_midpoint_head( + g, {lo_id: MIDPOINT + lower * FP32_ULP, hi_id: MIDPOINT + higher * FP32_ULP} + ) + expect = None if winner is None else [lo_id if winner == "lower" else hi_id] + _check_fused_against_stock(f"midpoint {where} {lower:+d}/{higher:+d} ulp", w, x, coded, expect) + + +# ------------------------------------------------------------------ the coded head against the BF16 fragments + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_coded_head_is_the_bf16_head_bit_for_bit(nreq) -> None: + """Integer data (fp32 sums exact in any order) with exceptions on both sides of the exponent window: the coded + head's winning keys (value and global id) and its bookkeeping are the BF16 fragments' bit for bit, on one rank and + across the ranks.""" + g = gen(seed_of(f"coded vs bf16 {nreq}")) + sign = torch.randint(0, 2, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * 2 - 1 + w = (torch.randint(1, 41, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * sign).to( + torch.bfloat16 + ) + # Exceptions on both sides of the window of 1..40: exact zeros, and 2^15 (fp32 sums stay below 2^24). + w[::5, ::512] = 0.0 + w[::13, 7::600] = 32768.0 + x = torch.randint(-2, 3, (nreq, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + wf_c, cd = _head(w, True) + assert cd[1].numel() > 0, "the shard must carry exceptions" + wf_b, cb = _head(w, False) + outs = [] + for wf, c in ((wf_c, cd), (wf_b, cb)): + args = _args(x, None, nreq, wf, 1, 1, 1, gen(seed_of("bookkeeping")), coded=c) + args = list(args) + args[9] = 0 + with validating(op): + op(*args, fixed_order=_fixed(args)) + torch.cuda.synchronize() + outs.append(args) + coded_args, bf16_args = outs + for i in (11, 13, 15, 16, 22): # draft_out, ids_next, pos_next, kv_lens, out_keys + assert torch.equal(coded_args[i], bf16_args[i]), i + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), DRAFT_SHARD) + assert torch.equal(coded_args[22], make_keys(best, idx)), ( + "the winning keys are not the in-tree ones" + ) + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_tiny_exception_products_are_exact(nreq) -> None: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1, and a one-hot + activation on column 0: the exceptions' exact products are the only non-zero logits, so the coded head must pick + row 1 with the BF16 head's key bit for bit (a correction against a non-zero placeholder cancels them to 0).""" + n, k = 256, 1024 + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + x = torch.zeros((nreq, k), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + wf_c, cd = _head(w, True) + assert cd[3] == 113 + keys = [] + for wf, c in ((wf_c, cd), _head(w, False)): + args = _args(x, None, nreq, wf, 0, 1, 1, gen(seed_of("tiny exceptions")), coded=c) + with validating(op): + # 16 tiles on an all-SM grid: every tile is streamed by all 16 warps, which only the fixed-order sum + # (D-FIX) adds deterministically (is_valid rejects the atomics there). + op(*args, fixed_order=True) + torch.cuda.synchronize() + assert args[13].tolist() == [1] * nreq + keys.append(args[22]) + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 0) + assert torch.equal(keys[0], keys[1]) and torch.equal(keys[0], make_keys(best, idx)) + + +def test_prepare_head_codes_the_shard_and_fails_closed(monkeypatch) -> None: + """With the coding on, prepare_head keeps the coded shard when it decodes to the weight bit for bit; a coding + that does not is dropped for the BF16 fragments. (The tp8 target ships HEAD_CODED = False.)""" + monkeypatch.setattr(draft_step_kernels, "HEAD_CODED", True) + g = gen(seed_of("prepare coded")) + w = randn((4096, 1024), g, 0.02) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 2) + assert prepared.coded and prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(head_weight(prepared.wf, 1024, prepared.exc_off, prepared.exc_pos, prepared.exc_w, + prepared.exp_base), w) # fmt: skip + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[11] ^= 0x10 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(draft_step_kernels, "encode_chunked", corrupt) + dropped = draft_step_kernels.prepare_head(head, 2) + assert not dropped.coded and dropped.exc_off is None + assert torch.equal(dropped.wf, bf16_fragments(w)) + + +def test_prepare_head_keeps_bf16_when_coding_is_off(monkeypatch) -> None: + """HEAD_CODED = False keeps the BF16 fragments even for a shard whose coding would decode bit for bit.""" + g = gen(seed_of("prepare bf16")) + w = randn((4096, 1024), g, 0.02) + monkeypatch.setattr(draft_step_kernels, "HEAD_CODED", False) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2) + assert not prepared.coded and prepared.exc_off is None + assert torch.equal(prepared.wf, bf16_fragments(w)) + + +def test_guard_refuses_head_rows_past_x() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[1] = args[1] + 1000 + with pytest.raises(AssertionError, match="x_rows must index x"): + with validating(op): + op(*args, fixed_order=_fixed(args)) + + +def test_guard_refuses_too_few_ctas() -> None: + args = list(_case(op.CELLS[1].spec, 2)) + args[20] = 8 + with pytest.raises(AssertionError, match="tiles per CTA"): + with validating(op): + op(*args, fixed_order=_fixed(args)) + + +# ------------------------------------------------- one pass as two launches (uGPU, kernels/m3_draft_step.UGPU) + + +@pytest.mark.parametrize( + "cell", + [c for c in op.CELLS if not c.spec["coded"]], + ids=[c.why[:44] for c in op.CELLS if not c.spec["coded"]], +) +def test_two_launches_per_domain_are_the_one_launch_pass(cell) -> None: + """A pass split into two launches (tiles [0, split) and [split, n), one per locality domain stream, the exchange + and bookkeeping by the last of both launches' CTAs) writes exactly what the one-launch pass writes: tokens, + winning keys, draft columns, next positions, KV and sequence lengths, and re-arms the counters and keys.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_locality import ( # noqa: E501 + LocalityMap, + ) + + one = _case(cell.spec, seed_of(cell.why)) + two = _case(cell.spec, seed_of(cell.why)) + op(*one) + lm = LocalityMap.get(torch.cuda.current_device()) + streams = lm.domain_streams() + n_tiles = two[3].shape[0] + split = max(1, n_tiles // 2) + per = min(lm.sms_per_domain, split, n_tiles - split) + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + for st, lo, hi in ((streams[0], 0, split), (streams[1], split, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + op(*two[:20], per, False, *two[22:], tile_lo=lo, tile_hi=hi, total_ctas=2 * per) + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) + torch.cuda.synchronize() + for i, name in ( + (11, "draft_out"), + (13, "ids_next"), + (15, "pos_next"), + (16, "kv_lens"), + (22, "out_keys"), + ): + assert torch.equal(one[i], two[i]), name + if one[18] is not None: + assert torch.equal(one[18], two[18]), "seq_lens" + assert two[6].tolist() == [0, 1] and not bool(two[5].any()), "counters / keys not re-armed" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py new file mode 100644 index 000000000000..601b65e50ed2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""D-FIX: the draft step's tile pieces and their summation order (kernels/m3_draft_step.FIXED_ORDER). + +A CTA's (tile, chunk) range is cut into 16 equal warp segments; the shared-memory atomics add a tile's pieces in +arrival order, which is deterministic only while no tile takes more than 2 pieces. ``fixed_order`` sums them in warp +order instead. The geometry tests enumerate the real grid (the shard on 212 SMs) so a grid or shard change cannot +silently bring back an order-dependent sum; the GPU tests hold the fixed order to the atomics on exact data and to the +warp-order oracle on a planted tile whose sum depends on the order. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( + m3_draft_step as kds, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_SHARD, + MAX_TOKENS, + bf16_fragments, + max_tile_pieces, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, head_workspace # noqa: E402 + +OP = "minimax_m3_sm107tp8_draft_step" +SMS = 212 # Rubin's SMs: the grid the draft step launches on +K = 6144 +CHUNKS = K // 16 +# This target's shard on Rubin: 7-8 tiles per CTA; a 7-tile CTA's tiles take up to 4 pieces (the atomics' sum is +# order-dependent, so the fixed order is on). +ORDER_DEPENDENT = True + +gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + + +@pytest.mark.parametrize("grid", [SMS, kds.GRID]) +def test_real_grid_pieces(grid) -> None: + """One CTA per SM and the DS-GRID launch (kernels/m3_draft_step.GRID).""" + pieces = max_tile_pieces(DRAFT_SHARD // 16, grid, CHUNKS) + assert (pieces > 2) == ORDER_DEPENDENT, (grid, pieces) + + +def test_ds_grid_launches_the_configured_grid(monkeypatch) -> None: + monkeypatch.setattr(kds, "_num_sms", lambda: SMS) + assert kds._grid() == min(kds.GRID, SMS) + monkeypatch.setattr(kds, "GRID", 0) + assert kds._grid() == SMS + + +def test_max_tile_pieces_counts_warp_segments() -> None: + # One tile per CTA: all 16 warps stream it. 7 tiles per CTA: a 168-chunk segment can sit inside a tile (TP8). + assert max_tile_pieces(16, 16, CHUNKS) == 16 + assert max_tile_pieces(7, 1, CHUNKS) == 4 + assert max_tile_pieces(8, 1, CHUNKS) == 2 + assert max_tile_pieces(15, 1, CHUNKS) == 2 + + +@pytest.mark.parametrize("flag", [True, False]) +def test_the_kernel_engages_the_fixed_order_where_the_geometry_needs_it(flag, monkeypatch) -> None: + monkeypatch.setattr(kds, "FIXED_ORDER", flag) + x = torch.empty(1, K) + + def head(n_tiles): + return types.SimpleNamespace(wf=torch.empty(n_tiles, 1), exc_off=None, bnd=None) + + assert kds._fixed_order(1, head(DRAFT_SHARD // 16), x, SMS) == (flag and ORDER_DEPENDENT) + assert kds._fixed_order(1, head(DRAFT_SHARD // 16), x, kds.GRID) == (flag and ORDER_DEPENDENT) + assert kds._fixed_order(1, head(16), x, 16) == flag # 16 pieces per tile + + +def _launcher(wf, x, fixed, nctas): + """(launch, ids, keys): one pass's op launch on buffers made here (outside any capture), world 1.""" + uc, keys, counters = head_workspace(1, slot_rows=MAX_TOKENS) + m = x.shape[0] + i32 = dict(dtype=torch.int32, device="cuda") + draft_out = torch.zeros(m, 4, **i32) + ids = torch.zeros(m, **i32) + pos_src, pos_next, kv_lens = ( + torch.full((m,), 100, **i32), + torch.zeros(m, **i32), + torch.full((m,), 500, **i32), + ) + out = torch.zeros(m, dtype=torch.int64, device="cuda") + + def launch(): + getattr(torch.ops.trtllm, OP)(x, None, m, wf, 0, keys, counters, uc, uc.data_ptr(), 0, 1, draft_out, 1, ids, + pos_src, pos_next, kv_lens, None, None, 3, nctas, False, out, + fixed_order=fixed) # fmt: skip + + return launch, ids, out + + +@gpu +@pytest.mark.parametrize("nctas,vs", [(SMS, DRAFT_SHARD), (16, 256)]) +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_fixed_order_equals_the_atomics_on_exact_data(nctas, vs, m) -> None: + """Integer data: every fp32 sum is exact in any order, so both summations give the same bits.""" + g = gen(17 + m) + w = torch.randint(-1, 2, (vs, K), generator=g, device="cuda").to(torch.bfloat16) + x = torch.randint(-1, 2, (m, K), generator=g, device="cuda").to(torch.bfloat16) + wf = bf16_fragments(w) + out = {} + for fixed in (False, True): + launch, ids, keys = _launcher(wf, x, fixed, nctas) + launch() + torch.cuda.synchronize() + out[fixed] = (keys, ids) + assert torch.equal(out[False][0], out[True][0]) + ref = (x.double() @ w.double().t()).float().bfloat16().float().argmax(dim=1).int() + assert torch.equal(out[True][1], ref) + + +@gpu +@pytest.mark.parametrize("graph", [False, True]) +def test_fixed_order_is_the_warp_order_sum_on_a_planted_tile(graph) -> None: + """One tile per CTA (16 pieces): row 0's pieces are +2^25 (warp 0), +1 (warps 1..14), -2^25 (warp 15). In warp + order the +1s are absorbed by 2^25 and the row sums to 0, so the competitor (row 255, logit 0.5) wins; arrival + order could give up to 14. Every launch must give the warp-order result.""" + w = torch.zeros(256, K, device="cuda") + seg = CHUNKS // 16 * 16 # one warp segment in K + w[0, 0] = 2.0**25 + for s in range(1, 15): + w[0, s * seg] = 1.0 + w[0, 15 * seg] = -(2.0**25) + w[255, 0] = 0.5 + wf = bf16_fragments(w.to(torch.bfloat16)) + x = torch.ones(1, K, device="cuda", dtype=torch.bfloat16) + launch, ids, _ = _launcher(wf, x, True, 16) + run = launch + if graph: + launch() + torch.cuda.synchronize() + gr = torch.cuda.CUDAGraph() + cs = torch.cuda.Stream() + with torch.cuda.stream(cs), torch.cuda.graph(gr, stream=cs): + launch() + run = gr.replay + for _ in range(16): + ids.zero_() + run() + torch.cuda.synchronize() + assert int(ids[0]) == 255 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py new file mode 100644 index 000000000000..e6d2fb864da9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py @@ -0,0 +1,376 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the forward-tail spec fold (kernels/m3_fold_tail and M3SpecStepFold.in_forward). + +The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers +(step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing +when the step was not armed. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( + m3_fold_tail, + m3_spec_tail, +) + +assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" + +VOCAB = 200064 +NUM_SLOTS = 9 +MAX_ROWS = 8 +STALE = -4 + + +def _store(widths, g): + wn, wx, wd = widths + i32 = dict(dtype=torch.int32, device="cuda") + return types.SimpleNamespace( + new_tokens=torch.randint( + 0, VOCAB, (wn, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_new_tokens=torch.randint( + 0, VOCAB, (wx, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_draft_tokens=torch.randint( + 0, VOCAB, (NUM_SLOTS, wd), generator=g, device="cuda", dtype=torch.int32 + ), + new_tokens_lens=torch.randint(1, 5, (NUM_SLOTS,), generator=g, **i32), + ) + + +def _clone(ns): + return types.SimpleNamespace(**{k: v.clone() for k, v in vars(ns).items()}) + + +def _outputs(rows, t, g): + d = t - 1 + return { + "new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_draft_tokens": torch.randint( + 0, VOCAB, (rows, d), generator=g, device="cuda", dtype=torch.int32 + ), + "new_tokens_lens": torch.randint( + 1, t + 1, (rows,), generator=g, device="cuda", dtype=torch.int32 + ), + } + + +def _engine(t): + i32 = dict(dtype=torch.int32, device="cuda") + return ( + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS * (t - 1),), STALE, **i32), + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS,), STALE, **i32), + ) + + +def _reference(store, engine, slots, out, t): + """slot_scatter into the store, then step_gather of the next step's inputs from it.""" + assert m3_spec_tail.slot_scatter( + slots, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["new_tokens_lens"], + store.new_tokens_lens, + out["next_draft_tokens"], + store.next_draft_tokens, + ) + p = slots.numel() + ids, draft, pos_off, kv_off = engine + slots32 = slots.to(torch.int32) + assert m3_spec_tail.step_gather( + store.next_new_tokens, + store.next_draft_tokens, + store.new_tokens_lens, + slots32, + slots32.repeat_interleave(t), + p, + t, + t - 1, + ids, + 0, + draft, + 0, + pos_off, + kv_off, + 0, + ) + + +def _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=False, one_warp=False): + ids, draft, pos_off, kv_off = engine + assert m3_fold_tail.fold_tail( + ctrl, + slots_buf, + rows, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["next_draft_tokens"], + store.next_draft_tokens, + out["new_tokens_lens"], + store.new_tokens_lens, + ids, + draft, + pos_off, + kv_off, + t, + cpp=cpp, + one_warp=one_warp, + ) + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("widths", [(4, 4, 3), (8, 8, 7)], ids=["store=step", "store wider"]) +def test_tail_matches_scatter_then_gather(rows, widths, cpp) -> None: + """Store bit for bit; engine rows [0, P*T) / [0, P) bit for bit (the forward reads only those rows).""" + t = 4 + g = torch.Generator(device="cuda").manual_seed(1000 * rows + widths[0]) + store = _store(widths, g) + out = _outputs(rows, t, g) + slots = torch.randperm(NUM_SLOTS, generator=g, device="cuda")[:rows] + ref_store, tail_store = _clone(store), _clone(store) + ref_engine, tail_engine = _engine(t), _engine(t) + _reference(ref_store, ref_engine, slots, out, t) + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + slots_buf = torch.zeros(MAX_ROWS, dtype=torch.long, device="cuda") + slots_buf[:rows] = slots + _tail(tail_store, tail_engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(ref_store, name), getattr(tail_store, name)), name + spans = (rows * t, rows * (t - 1), rows * t, rows) + for k, (ref, got, n) in enumerate(zip(ref_engine, tail_engine, spans)): + assert torch.equal(ref[:n], got[:n]), k + assert torch.all(got[n:] == STALE), k + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +def test_unarmed_tail_writes_nothing(cpp) -> None: + """ctrl[0] = 0 (the step was not armed): the store and the engine buffers are untouched.""" + t, rows = 4, 2 + g = torch.Generator(device="cuda").manual_seed(7) + store = _store((4, 4, 3), g) + out = _outputs(rows, t, g) + before, engine = _clone(store), _engine(t) + ctrl = torch.zeros(1, dtype=torch.int32, device="cuda") + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(before, name), getattr(store, name)), name + for buf in engine: + assert torch.all(buf == STALE) + + +def _fold(): + fold = modeling.M3SpecStepFold() + t = 4 + engine = _engine(t) + fold.bind_engine_buffers(*engine, max_batch_size=MAX_ROWS) + g = torch.Generator(device="cuda").manual_seed(11) + fold.bind_sampler_store(_store((4, 4, 3), g)) + return fold, g, t + + +def test_fold_handshake_stages_only_confirmed_rows() -> None: + """begin_step arms; the sampler confirms only the armed rows from the tail's outputs; the next step is staged only + for exactly those rows in that order.""" + fold, g, t = _fold() + assert fold.begin_step([3, 5], 2) is False # nothing staged before the first forward + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [3, 5]) is True + assert fold.begin_step([3, 5], 2) is True + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert ( + fold.sampler_rows_stored(out, 0, [5, 3]) is False + ) # other order: the sampler stores itself + assert fold.begin_step([5, 3], 2) is False + + +def test_fold_handshake_needs_the_tail_outputs() -> None: + """Outputs the tail did not run on (another dict, a mixed step, an unarmed step) are never reported stored.""" + fold, g, t = _fold() + fold.begin_step([1], 1) + fold.store_in_forward(_outputs(1, t, g)) + assert fold.sampler_rows_stored(_outputs(1, t, g), 0, [1]) is False + fold.begin_step(None, 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [1]) is False + fold.begin_step([1], 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 1, [1]) is False # rows were skipped (chunked context) + + +def test_a_padded_step_after_an_armed_one_writes_nothing() -> None: + """K3 #150 check #5: CUDA-graph padding rows carry no seq slot, so the engine passes rows=None for a padded step. + Even right after an armed step, that step uploads ctrl 0 and its forward tail leaves the store and the engine + buffers untouched (a dummy row never writes a live slot).""" + fold, g, t = _fold() + fold.begin_step([2, 4], 2) + fold.store_in_forward(_outputs(2, t, g)) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 2 + + fold.begin_step(None, 4) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 0 + store_before = _clone(fold.store) + engine_before = [buf.clone() for buf in fold.step_inputs] + fold.store_in_forward(_outputs(4, t, g)) + torch.cuda.synchronize() + for name in vars(store_before): + assert torch.equal(getattr(store_before, name), getattr(fold.store, name)), name + for before, buf in zip(engine_before, fold.step_inputs): + assert torch.equal(before, buf) + + +def test_the_fold_stores_in_the_forward_once_bound() -> None: + """The fold always stores and stages at the forward's tail; before the engine binds its buffers nothing is armed.""" + fold = modeling.M3SpecStepFold() + assert fold.in_forward + assert fold.begin_step([0], 1) is False + + +def test_the_store_copy_event_is_recorded_at_binding(): + """An external event the sampler records after its side-stream store copies and the forward's tail waits on, + recorded once at binding so a captured wait has a completed record to wait on.""" + fold = modeling.M3SpecStepFold() + device = torch.device("cuda") + fold.bind_engine_buffers( + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + 4, + ) + assert fold.store_copies_done is not None + fold.store_copies_done.synchronize() + + +def test_the_cpp_tail_launches_plainly() -> None: + """The C++ tail writes the next step's inputs, which pre-wait readers take before their grid wait (MSA contract): + its kernel node carries no programmatic (PDL) serialization.""" + from cuda.bindings import runtime as rt + + t, rows = 4, 2 + g = torch.Generator(device="cuda").manual_seed(13) + store, out, engine = _store((4, 4, 3), g), _outputs(rows, t, g), _engine(t) + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=True) + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + # cuda-python changed the signature: (graph[, numNodes]) on older bindings, (graph, nodes, numNodes) on newer ones. + try: + err, _, n = rt.cudaGraphGetNodes(raw) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + except TypeError: + err, _, n = rt.cudaGraphGetNodes(raw, None, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, None, n) + flags = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind == rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + # Older bindings return no attribute for a kernel launched without one (a plain launch). + flags.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + assert flags == [0], flags + + +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("widths", [(4, 4, 3), (8, 8, 7)], ids=["store=step", "store wider"]) +def test_one_warp_form_is_bitwise_the_kernel(rows, widths) -> None: + """The in-kernel form (one warp over the rows: what the drafter's last-pass draft_step runs in its last CTA, H-R2-3) + writes exactly what the one-CTA-per-row kernel writes, armed and unarmed, eager and from a graph.""" + t = 4 + for trial in range(3): + g = torch.Generator(device="cuda").manual_seed(100 * rows + trial) + store, out, engine = _store(widths, g), _outputs(rows, t, g), _engine(t) + slots_buf = torch.randperm(NUM_SLOTS, generator=g, device="cuda")[:MAX_ROWS].to(torch.long) + for armed in (rows, 0): + ctrl = torch.tensor([armed], dtype=torch.int32, device="cuda") + a_store, a_engine = _clone(store), tuple(e.clone() for e in engine) + b_store, b_engine = _clone(store), tuple(e.clone() for e in engine) + _tail(a_store, a_engine, ctrl, slots_buf, out, t, rows, cpp=True) + _tail(b_store, b_engine, ctrl, slots_buf, out, t, rows, cpp=True, one_warp=True) + torch.cuda.synchronize() + for k in vars(a_store): + assert torch.equal(getattr(a_store, k), getattr(b_store, k)), (k, armed, trial) + for x, y in zip(a_engine, b_engine): + assert torch.equal(x, y), (armed, trial) + # From a captured graph, replayed on fresh outputs. + g = torch.Generator(device="cuda").manual_seed(7) + store, out, engine = _store(widths, g), _outputs(rows, t, g), _engine(t) + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + b_store, b_engine = _clone(store), tuple(e.clone() for e in engine) + _tail(b_store, b_engine, ctrl, slots_buf, out, t, rows, cpp=True, one_warp=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _tail(b_store, b_engine, ctrl, slots_buf, out, t, rows, cpp=True, one_warp=True) + for step in range(3): + fresh = _outputs(rows, t, torch.Generator(device="cuda").manual_seed(20 + step)) + for k in out: + out[k].copy_(fresh[k]) + a_store, a_engine = _clone(b_store), tuple(e.clone() for e in b_engine) + _tail(a_store, a_engine, ctrl, slots_buf, out, t, rows, cpp=True) + graph.replay() + torch.cuda.synchronize() + for k in vars(a_store): + assert torch.equal(getattr(a_store, k), getattr(b_store, k)), (k, step) + for x, y in zip(a_engine, b_engine): + assert torch.equal(x, y), step + + +def test_the_store_wait_precedes_only_a_plain_launch_under_capture(): + """The store wait is an external-event wait node under graph capture, which drops the next kernel's PDL edge: + it may precede the fold tail (a plain launch); any other launch fails the capture (eager steps are not checked).""" + fold = modeling.M3SpecStepFold() + device = torch.device("cuda") + fold.bind_engine_buffers( + *(torch.zeros(8, dtype=torch.int32, device=device) for _ in range(4)), 4 + ) + assert fold.STORE_WAIT_BEFORE == frozenset({"fold_tail"}) + fold._wait_store_copies(before="draft_step") # eager: allowed + x = torch.zeros(4, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + fold._wait_store_copies(before="fold_tail") + x.add_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(x, torch.ones(4, device=device)) + bad = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="PDL edge"): + with torch.cuda.graph(bad, stream=stream): + fold._wait_store_copies(before="draft_step") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py new file mode 100644 index 000000000000..0f6b4f255102 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the per-request forced-acceptance draw (``FORCE_AL_PER_REQUEST``, the ``m3_force_vec_req`` entry, +synth only): each request's forced count comes from its own key (its KV length at the verify), so it does not depend +on the call history, on the request's slot in the batch, or on the other requests.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec_req, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SEED, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _pool(): + g = torch.Generator(device="cpu") + g.manual_seed(_FORCE_ACCEPT_RNG_SEED) + return torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, dtype=torch.float32, generator=g).cuda() + + +def _draw(pool, keys, base_total=2, max_total=4, frac=0.6): + out = torch.empty(keys.numel(), dtype=torch.int32, device="cuda") + m3_force_vec_req( + pool, + keys, + out, + keys.numel(), + base_total, + max_total, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (keys.numel() - 1).bit_length()), + ) + return out + + +def test_a_request_draws_the_same_whatever_ran_before_or_beside_it() -> None: + pool = _pool() + keys = torch.tensor([165_521, 40_003, 323_007, 7], dtype=torch.int32, device="cuda") + first = _draw(pool, keys) + for _ in range(5): # no shared counter: repeated calls draw the same + assert torch.equal(_draw(pool, keys), first) + perm = torch.tensor([2, 0, 3, 1], device="cuda") + assert torch.equal(_draw(pool, keys[perm]), first[perm]) # slot-independent + alone = torch.cat([_draw(pool, keys[i : i + 1]) for i in range(4)]) + assert torch.equal(alone, first) # independent of the other requests + + +def test_the_counts_follow_the_forced_fraction() -> None: + pool = _pool() + keys = torch.arange(100_000, 100_000 + 64 * 400, dtype=torch.int32, device="cuda") + counts = torch.cat([_draw(pool, keys[i : i + 64]) for i in range(0, keys.numel(), 64)]) + assert set(counts.tolist()) <= {2, 3} + extra = (counts == 3).float().mean().item() + assert abs(extra - 0.6) < 0.02, extra diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..6b63ba265112 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec entry of the MiniMax-M3 tp8 target: the fractional forced-acceptance draw.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(5) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + counter = torch.tensor([41], dtype=torch.int64, device="cuda") + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + counter, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert int(counter.item()) == 42, "the draw bumps the step counter once" + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py new file mode 100644 index 000000000000..871dfa22d617 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec_req entry of the MiniMax-M3 tp8 target: the per-request fractional forced-acceptance +draw (one draw per request at its own key, no shared counter).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec_req as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(7) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + # Each request's key is its KV length at the verify: long, distinct, arbitrary. + keys = torch.randint( + 1, 1 << 20, (spec["batch"],), generator=g, device="cuda", dtype=torch.int64 + ).to(torch.int32) + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + keys, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py new file mode 100644 index 000000000000..b66a1e35e4a4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (kernels/m3_lowm_splitk, +modeling.DRAFT_LOWM_OWN_SPLITK): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's +linears, eager and in a CUDA graph, with and without PDL.""" + +import pytest +import torch +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, + run_splitk_dense_swiglu, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + TACTICS, +) + +assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" + +# Every split-K cell of the TP8 drafter table: (N, K, max rows, tactic) for kind "s". +CELLS = [ + (nk[0], nk[1], rows, tactic) + for nk, entries in sorted(TACTICS.items()) + for rows, kind, tactic in entries + if kind == "s" +] + + +def test_the_table_has_the_four_certified_split_k_cells() -> None: + assert sorted((n, k, m) for n, k, m, _ in CELLS) == [ + (6144, 1024, 1), + (6144, 1024, 4), + (6144, 2304, 1), + (6144, 2304, 4), + ] + + +@pytest.mark.parametrize("use_pdl", [False, True], ids=["plain", "pdl"]) +@pytest.mark.parametrize("cell", CELLS, ids=[f"{n}x{k}@{m}" for n, k, m, _ in CELLS]) +def test_bitwise_flashinfer(cell, use_pdl) -> None: + n, k, m, tactic = cell + g = torch.Generator(device="cuda").manual_seed(n + 7 * k + m) + for trial in range(3): + x = (torch.randn(m, k, generator=g, device="cuda") * (1.0 + trial)).to(torch.bfloat16) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + ref = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + got = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + fi_run(x, w.t(), None, ref, use_pdl, FiTactic(*tactic)) + run_splitk_dense(x, w.t(), None, got, use_pdl, SplitKTactic(*tactic)) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), f"trial {trial}" + # In a CUDA graph, as the drafter runs it. + cap = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_splitk_dense(x, w.t(), None, cap, use_pdl, SplitKTactic(*tactic)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(cap.view(torch.int16), ref.view(torch.int16)) + + +# The drafter MLP's down projection: the SwiGLU folded into the GEMM's B load (c2). +DOWN_CELLS = [c for c in CELLS if (c[0], c[1]) == (6144, 2304)] + + +@pytest.mark.parametrize("use_pdl", [False, True], ids=["plain", "pdl"]) +@pytest.mark.parametrize("cell", DOWN_CELLS, ids=[f"{n}x{k}@{m}" for n, k, m, _ in DOWN_CELLS]) +def test_swiglu_prologue_is_bitwise_silu_then_split_k(cell, use_pdl) -> None: + n, k, m, tactic = cell + g = torch.Generator(device="cuda").manual_seed(31 * n + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + scratch = torch.full((8, k), 7, dtype=torch.bfloat16, device="cuda") + for trial in range(3): + gate_up = (torch.randn(m, 2 * k, generator=g, device="cuda") * (2.0 + trial)).to( + torch.bfloat16 + ) + gate_up[:, :6] = torch.tensor([60.0, -60.0, 0.0, -0.0, 88.0, -88.0], dtype=torch.bfloat16) + act = torch.ops.trtllm.silu_and_mul(gate_up) + ref = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + fi_run(act, w.t(), None, ref, use_pdl, FiTactic(*tactic)) + got = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + run_splitk_dense_swiglu(gate_up, w.t(), got, use_pdl, SplitKTactic(*tactic), scratch) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), f"trial {trial}" + assert torch.equal(scratch[:m].view(torch.int16), act.view(torch.int16)), ( + "the B scratch holds the activation" + ) + cap = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_splitk_dense_swiglu(gate_up, w.t(), cap, use_pdl, SplitKTactic(*tactic), scratch) + scratch.fill_(3) # a replay rewrites what it reads + graph.replay() + torch.cuda.synchronize() + assert torch.equal(cap.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..f4739cf8f5fc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp8 m3_sampled_accept entry (fused temperature + top-p verify over the 25008-column TP8 shard) +and its workspace size.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_sampled_accept import ( # noqa: E501 + m3_sampled_accept as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_sampled_accept_ws_words import ( # noqa: E501 + m3_sampled_accept_ws_words as ws_words, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" + +DRAFT_LEN = 3 +#: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). +HEAD = 12 +#: Draws per distribution cell, and the total-variation band against the reference distribution. With 12 +#: categories and >= 4000 draws the expected TV of an exact sampler is ~0.02 (sd ~0.006); 0.06 is ~6 sd above it +#: while a wrong nucleus cut (one category too many or too few, >= 1 / 12 of the mass moved) lands at >= 0.08. +DRAWS = 4000 +TV_BAND = 0.06 + + +def _logits(spec, g, rows): + vs = spec["vs"] + if spec["check"] == "flat": + return (torch.randn((rows, vs), generator=g, device="cuda") * 0.3).to(torch.bfloat16) + if spec["check"] == "argmax": + return (torch.randn((rows, vs), generator=g, device="cuda") * 3).to(torch.bfloat16) + row = torch.full((vs,), -20.0, device="cuda") + hot = torch.randperm(vs, generator=g, device="cuda")[:HEAD] + row[hot] = torch.linspace(2.0, 0.2, HEAD, device="cuda") + return row.to(torch.bfloat16).expand(rows, vs).contiguous() + + +def _case(spec, seed): + g = gen(seed) + nreq = spec["nreq"] + rows = nreq * (DRAFT_LEN + 1) + logits = _logits(spec, g, rows) + words = ws_words(1) + uc = torch.zeros(words, dtype=torch.int64, device="cuda") + counters = torch.zeros(2, dtype=torch.int32, device="cuda") + tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + extra = torch.zeros(1, dtype=torch.int32, device="cuda") + temperatures = torch.full((rows,), spec["temperature"], dtype=torch.float32, device="cuda") + top_ps = torch.full((rows,), spec["top_p"], dtype=torch.float32, device="cuda") + seeds = torch.randint(0, 2**62, (rows,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.zeros(rows, dtype=torch.int64, device="cuda") + # Drafts drawn from the likeliest tokens, so some prefixes match. + top = logits[0].float().topk(3).indices.to(torch.int32) + spec["vocab_offset"] + draft = top[torch.randint(0, 3, (nreq, DRAFT_LEN), generator=g, device="cuda")].contiguous() + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + force = spec.get("force", "none") + force_total = 3 if force == "total" else 0 + force_vec = (torch.randint(1, DRAFT_LEN + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if force == "vec" else None) # fmt: skip + return [logits, spec["vocab_offset"], temperatures, top_ps, seeds, offsets, uc, uc.data_ptr(), counters, tokens, + 0, 1, draft, batch_idx, acc, num, gather, force_total, next_new, force_vec, extra, True] # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Statistical for the draw, exact for everything derived from it. + + Every drawn token has nonzero reference probability (lies in the nucleus); the acceptance outputs equal the + reference's strict acceptance of the drawn tokens, every call; the call counter advances once per call; and, + for the distribution cells, the draws over many calls (fresh Philox offsets) are within TV_BAND of the + reference distribution (see DRAWS / TV_BAND). + """ + spec = cell.spec + args = _case(spec, seed_of(cell.why)) + logits, offset, rows = args[0], spec["vocab_offset"], args[0].shape[0] + certifying(op) + ref = op.reference(*args) + probs = ref["probs"] + calls = 1 if spec["check"] == "argmax" else -(-DRAWS // rows) + counts = torch.zeros(logits.shape[1], dtype=torch.float64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + drawn = args[9][:rows].clone() + local = drawn.long() - offset + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + assert bool((probs.gather(1, local[:, None]) > 0).all()), ( + "a drawn token outside the nucleus" + ) + acc = ref["accept"](drawn) + assert torch.equal(args[14], acc["accepted"]) + assert torch.equal(args[15], acc["num_accepted"]) + assert torch.equal(args[16], acc["gather_ids"]) + assert torch.equal(args[18], acc["next_new"]) + assert args[8].tolist() == [0, c + 1] + counts.index_add_(0, local, torch.ones(rows, dtype=torch.float64, device="cuda")) + if spec["check"] == "distribution": + emp = counts / counts.sum() + tv = 0.5 * (emp - probs[0].double()).abs().sum().item() + assert tv <= TV_BAND, f"total variation {tv:.3f} against the reference nucleus distribution" + + +def _walk_outputs(args, one_round_k, ilp, calls=3, offset_step=1000): + """The draws and acceptance outputs of `calls` calls with the owner walk one token (ilp False) or 8 tokens at a + time (ilp True), at SA-1R variant `one_round_k`, each call at its own Philox offsets.""" + if one_round_k: + words = torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words(1, one_round_k) + args[6] = torch.zeros(words, dtype=torch.int64, device="cuda") + args[7] = args[6].data_ptr() + rows = args[0].shape[0] + got = [] + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * offset_step) + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept( + *args, True, False, one_round_k, None, ilp + ) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + return got + + +@pytest.mark.parametrize("one_round_k", [0, 82]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_ilp_walk_changes_nothing(cell, one_round_k) -> None: + """kernels/m3_sampled_accept.ILP_WALK (SA-PW) only computes the owner's token weights 8 at a time; the walk adds + them in the same order with the same early exit, so the draws and acceptance outputs are the one-token walk's, + bit for bit, call after call, with SA-1R off and on.""" + one = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, False) + ilp = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, True) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.95]) +def test_the_ilp_walk_changes_nothing_over_many_rounds(top_p) -> None: + """A flat tail (many tokens of equal weight) under a tiny and a wide nucleus: rows take many exchange rounds and + each round's walk exits at a different token; the 8-token walk draws the same tokens as the one-token walk.""" + spec = dict(nreq=1, vs=25008, vocab_offset=25008, check="flat", top_p=top_p, temperature=1.0) + one = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), 0, False, calls=24) + ilp = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), 0, True, calls=24) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +def test_the_ilp_walk_is_the_same_under_graph_replay() -> None: + """The 8-token walk captured in a CUDA graph and replayed draws what the one-token walk draws eagerly.""" + cell = op.CELLS[0] + eager = _walk_outputs(_case(cell.spec, seed_of(cell.why)), 0, False) + _walk_outputs( + _case(cell.spec, seed_of("ilp walk warm-up")), 0, True, calls=1 + ) # first launch outside the capture + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept(*args, True, False, 0, None, True) + torch.cuda.current_stream().wait_stream(stream) + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + graph.replay() + torch.cuda.synchronize() + for x, y in zip(eager[c], [args[i] for i in (9, 14, 15, 16, 18)]): + assert torch.equal(x, y), ( + f"replay {c}: the graphed 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.0002]) +def test_long_rejection_chains_stay_in_the_nucleus(top_p) -> None: + """A tiny nucleus over a flat tail: rows that need more exchange rounds than the workspace has slots. + + Each round's candidate is drawn over the tokens heavier than the last one, so a row takes ~ln(1 / top_p) + 1 + rounds on average, with a tail past the 16 exchange slots (about 1% of the rows at top_p 0.001). Rounds reuse the + slots cyclically, told apart by a lap-parity bit in the wire tags, and no round accepts unconditionally: every + draw, however many rounds it took, lies in the reference nucleus. World 1 (one rank's 8 CTAs per row exchange + through the same slot ring); the cross-rank path needs the multi-device job. + """ + spec = dict(nreq=1, vs=25008, vocab_offset=25008, check="flat", top_p=top_p, temperature=1.0) + args = _case(spec, seed_of(f"flat tail top_p {top_p}")) + logits, rows = args[0], args[0].shape[0] + probs = op.reference(*args)["probs"] + calls = -(-DRAWS // rows) + outside = torch.zeros((), dtype=torch.int64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + local = args[9][:rows].long() - spec["vocab_offset"] + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + outside += (probs.gather(1, local[:, None]) == 0).sum() + assert args[8].tolist() == [0, c + 1] + draws = calls * rows + rounds = 2.0 + args[20].item() / draws + kept = int((probs > 0).sum(1).min().item()) + assert outside.item() == 0, ( + f"{outside.item()} of {draws} draws outside the nucleus (kept >= {kept} of {logits.shape[1]}, " + f"{rounds:.2f} rounds per row on average)" + ) + assert rounds > 4.0, ( + f"only {rounds:.2f} rounds per row: the rows no longer stress the slot ring" + ) + + +@pytest.mark.parametrize("cell", ws_words.CELLS, ids=[c.why[:44] for c in ws_words.CELLS]) +def test_workspace_words(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(ws_words) + world = cell.spec["world"] + assert ws_words(world) == ws_words.reference(world) + + +def test_guard_refuses_top_k_rows() -> None: + args = _case(op.CELLS[0].spec, 1) + top_ks = torch.full((args[0].shape[0],), 20, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="top-k rows"): + with validating(op): + op(*args, top_ks=top_ks) + + +def test_guard_refuses_fp64_temperatures() -> None: + args = _case(op.CELLS[0].spec, 2) + args[2] = args[2].double() + with pytest.raises(AssertionError, match="fp32 sampling parameters"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize( + "check,top_p", [("flat", 0.95), ("argmax", 0.95), ("flat", 0.001), ("distribution", 0.5)] +) +def test_the_copy_is_the_tp2_kernel_bit_for_bit(check, top_p) -> None: + """This target's own op (trtllm::minimax_m3_sm107tp8_sampled_accept) is the TP2 kernel it copies, bit for bit: + the same tokens, acceptance outputs, call counter and extra-round count from the same seeds and offsets.""" + tp2 = getattr(torch.ops.trtllm, "minimax_m3_sm107tp2_sampled_accept", None) + if tp2 is None: + pytest.skip("the TP2 op is not loaded") + spec = dict(nreq=1, vs=25008, vocab_offset=25008, check=check, top_p=top_p, temperature=1.0) + for call in range(8): + outs = [] + for which in ("tp8", "tp2"): + args = _case(spec, seed_of(f"copy {check} {top_p}")) + args[5].copy_( + torch.arange(args[0].shape[0], device="cuda", dtype=torch.int64) + call * 1000 + ) + if which == "tp8": + op(*args) + else: + tp2(*args[:21], args[21]) + torch.cuda.synchronize() + outs.append([args[i].clone() for i in (9, 14, 15, 16, 18, 20)] + [args[8].clone()]) + for a, b in zip(*outs): + assert torch.equal(a, b) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py new file mode 100644 index 000000000000..914e0d38bf05 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SA-1R (kernels/m3_sampled_accept.ONE_ROUND_K) on trtllm::minimax_m3_sm107tp8_sampled_accept, world 1, the TP8 shard. + +The one-round variants (one_round_k 9: exact top-9 list; 82: pooled top-8 + bound) must return the two-round +kernel's (one_round_k 0) sampled tokens, accepted tokens, accepted counts, gather ids and extra-attempt count BIT FOR +BIT on the same Philox stream, over random AgentX-like rows (several temperatures / top-p, broad and heavy-tailed +heads) and over rows built so the decision sits within ulps of top_p * Z (the margin fallback) or a slice holds more +heavy tokens than its list (the incomplete-list fallback), eager and under CUDA-graph replay. The round-0 decided +counter must engage on the AgentX-like rows and stay below the row count on the boundary rows (fallbacks fire). +Ported from MoE's prototype test (moe/sa1r/6c753632ec, exact on 39,968 rows, 788162). +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + +VS, D = 25008, 3 +KX = (9, 82) + +if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the sm_107 tp8 sampled accept runs on SM 10.7", allow_module_level=True) + + +class Runner: + """One variant with its own workspace / counters (call tag sequence) and outputs.""" + + def __init__(self, k, rows): + self.k, self.rows = k, rows + words = int(torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words(1, k)) + self.ws = torch.zeros(words, dtype=torch.int64, device="cuda") + self.counters = torch.zeros(2, dtype=torch.int32, device="cuda") + self.tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + self.extra = torch.zeros(1, dtype=torch.int32, device="cuda") + self.decided = torch.zeros(1, dtype=torch.int32, device="cuda") + nreq = rows // (D + 1) + self.acc = torch.zeros((nreq, D + 1), dtype=torch.int32, device="cuda") + self.nacc = torch.zeros(nreq, dtype=torch.int32, device="cuda") + self.gids = torch.zeros(nreq, dtype=torch.int64, device="cuda") + + def __call__(self, lg, temps, tps, seeds, offs, draft, bidx): + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept( + lg, 0, temps, tps, seeds, offs, self.ws, self.ws.data_ptr(), self.counters, self.tokens, 0, 1, draft, + bidx, self.acc, self.nacc, self.gids, 0, None, None, self.extra, False, True, False, self.k, + self.decided) # fmt: skip + + def out(self): + return (self.tokens[: self.rows].clone(), self.acc.clone(), self.nacc.clone(), self.gids.clone(), + self.extra.clone()) # fmt: skip + + +def _params(rows, temp, top_p, seed0): + temps = torch.full((rows,), temp, dtype=torch.float32, device="cuda") + tps = ( + (top_p if torch.is_tensor(top_p) else torch.full((rows,), top_p)) + .float() + .cuda() + .contiguous() + ) + seeds = torch.arange(rows, dtype=torch.int64, device="cuda") * 7919 + seed0 + offs = torch.zeros(rows, dtype=torch.int64, device="cuda") + return temps, tps, seeds, offs + + +def _sets(n, rows, g, kind): + out = [] + for _ in range(n): + lg = torch.randn(rows, VS, device="cuda", generator=g) * (1.0 if kind == "flat" else 2.0) + if kind == "lm": # a few heavy tokens over a broad tail + hot = torch.randint(0, VS, (rows, 8), device="cuda", generator=g) + lg.scatter_(1, hot, 12.0 + 3.0 * torch.rand(rows, 8, device="cuda", generator=g)) + elif kind == "tail": # 64 near-equal top tokens: lists overflow + hot = torch.randint(0, VS, (rows, 64), device="cuda", generator=g) + lg.scatter_(1, hot, 9.0 + 0.5 * torch.rand(rows, 64, device="cuda", generator=g)) + out.append(lg.to(torch.bfloat16).contiguous()) + return out + + +def _draft_from(lg, g): + nreq = lg.shape[0] // (D + 1) + am = lg.float().argmax(1).view(nreq, D + 1)[:, :D].int() + rnd = torch.randint(0, VS, (nreq, D), device="cuda", generator=g, dtype=torch.int32) + return torch.where(torch.rand(nreq, D, device="cuda", generator=g) < 0.7, am, rnd).contiguous() + + +def _compare(ref, got, what): + for n, x, y in zip( + ("tokens", "acc_tokens", "num_acc", "gather_ids", "extra_attempts"), ref, got + ): + assert torch.equal(x, y), f"{what}: {n} differs" + + +CASES = [(1.0, 0.95, "lm"), (0.7, 0.8, "lm"), (1.3, 0.5, "lm"), (1.0, 1.0, "lm"), (1.0, 0.95, "flat"), + (1.0, 0.95, "tail"), (0.6, 0.99, "tail")] # fmt: skip + + +def test_one_round_is_two_round_bitwise(): + g = torch.Generator(device="cuda").manual_seed(2) + total = 0 + decided = {k: 0 for k in KX} + for rows, calls in ((32, 20), (4, 60)): + for ci, (t, p, kind) in enumerate(CASES): + runners = {k: Runner(k, rows) for k in (0, *KX)} + temps, tps, seeds, offs = _params(rows, t, p, 1000 + ci) + sets = _sets(4, rows, g, kind) + bidx = torch.arange(rows // (D + 1), dtype=torch.int32, device="cuda") + for i in range(calls): + lg = sets[i % 4] + draft = _draft_from(lg, g) + offs.fill_(i) + for r in runners.values(): + r(lg, temps, tps, seeds, offs, draft, bidx) + for k in KX: + _compare( + runners[0].out(), runners[k].out(), f"K={k} rows={rows} case={ci} call={i}" + ) + total += rows + if kind == "lm" and p < 1.0: + for k in KX: + decided[k] += int(runners[k].decided.item()) + assert int(runners[0].decided.item()) == 0 + assert total >= 5000 + for k in KX: + assert decided[k] > 0, ( + f"K={k}: no row decided in round 0 on the AgentX-like rows (SA-1R not engaged)" + ) + + +def _boundary_rows(rows, n_heavy_same_slice): + """Heavy tokens heavier than a 50-token candidate group; top_p swept across ulps of S / Z (margin fallback), or + n heavy tokens in one slice (more than its list: incomplete fallback).""" + lg = torch.full((rows, VS), -40.0, device="cuda") + sl = ((VS + 7) // 8 + 7) // 8 * 8 + heavy = list(range(5, 5 + n_heavy_same_slice)) if n_heavy_same_slice else [5, 3 * sl + 11] + grp = [2 * sl + 3 + 37 * j for j in range(25)] + [5 * sl + 7 + 41 * j for j in range(25)] + lg[:, heavy] = 5.0 + lg[:, grp] = 4.0 + lg = lg.to(torch.bfloat16).contiguous() + w = torch.exp(lg[0].double() - lg[0].double().max()) + crit = float(w[heavy].sum()) / float(w.sum()) + ks = torch.arange(rows, dtype=torch.float64) - rows // 2 + return lg, (crit * (1.0 + ks * 2.0**-24)).float() + + +@pytest.mark.parametrize("graph", [False, True]) +def test_one_round_boundary_and_incomplete(graph): + rows = 32 + g = torch.Generator(device="cuda").manual_seed(3) + seen = {k: [0, 0] for k in KX} # decided rows, rows + for n_same in (0, 12): + lg, tps_rows = _boundary_rows(rows, n_same) + for shift in range(-4, 4): + tps = (tps_rows.double() * (1.0 + shift * 32 * 2.0**-24)).float() + runners = {k: Runner(k, rows) for k in (0, *KX)} + temps, tp_t, seeds, offs = _params(rows, 1.0, tps, 5000 + shift) + bidx = torch.arange(rows // (D + 1), dtype=torch.int32, device="cuda") + draft = _draft_from(lg, g) + graphs = {} + if graph: + for r in runners.values(): + r(lg, temps, tp_t, seeds, offs, draft, bidx) + torch.cuda.synchronize() + for k, r in runners.items(): + graphs[k] = torch.cuda.CUDAGraph() + with torch.cuda.graph(graphs[k]): + r(lg, temps, tp_t, seeds, offs, draft, bidx) + for r in runners.values(): + r.decided.zero_() + for i in range(12): + offs.fill_(i) + for k, r in runners.items(): + if graph: + graphs[k].replay() + else: + r(lg, temps, tp_t, seeds, offs, draft, bidx) + for k in KX: + _compare(runners[0].out(), runners[k].out(), + f"boundary n_same={n_same} shift={shift} call={i} K={k}") # fmt: skip + for k in KX: + seen[k][0] += int(runners[k].decided.item()) + seen[k][1] += 12 * rows + for k in KX: + assert 0 < seen[k][0] < seen[k][1], ( + f"K={k}: decided {seen[k]}: the boundary sweep missed both sides" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..ca795ef4f3ca --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_shard_accept entry (greedy verify from this rank's logits shard).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + craft_drafts, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" + + +def _case(spec, seed, edit_logits=None): + g = gen(seed) + nreq, d, vs, world, rank = ( + spec["nreq"], + spec["draft_len"], + spec["vs"], + spec["world"], + spec["rank"], + ) + m = nreq * (d + 1) + if spec["strided"]: + logits = randn((m, vs + 72), g, 3.0)[:, :vs] + else: + logits = randn((m, vs), g, 3.0) + if edit_logits is not None: + edit_logits(logits) + uc, keys, counters = head_workspace(world, slot_rows=SHARD_ROWS) + own = logits.float().max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=SHARD_ROWS, ties=spec.get("ties", False) + ) + offset = rank * vs + best, idx = local_argmax(logits.float(), offset) + tokens = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, SHARD_ROWS) + draft = craft_drafts(tokens, nreq, d) + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + force_total = 3 if spec["force"] == "total" else 0 + force_vec = (torch.randint(1, d + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if spec["force"] == "vec" else None) # fmt: skip + return (logits, offset, keys, counters, uc, uc.data_ptr(), rank, world, draft, batch_idx, acc, num, gather, + force_total, next_new, num_sms(), force_vec) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bit-exact: the op reduces BF16 logits without arithmetic, ties (within the shard and across ranks) go to + the lowest global index on both sides.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) + assert args[3].tolist() == [0, 1] and not bool(args[2].any()) + + +def test_guard_refuses_nan_logits() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0][1, 7] = float("nan") + with pytest.raises(AssertionError, match="NaN"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_batch_index_past_the_requests() -> None: + args = list(_case(op.CELLS[0].spec, 2)) + args[9] = args[9] + 1 + with pytest.raises(AssertionError, match="batch_idx"): + with validating(op): + op(*args) + + +def _signed_zero_maxima(logits: torch.Tensor) -> None: + """Every row's maximum is zero, reached as -0.0 at a lower index and +0.0 at a higher one.""" + logits.fill_(-1.0) + logits[:, 3] = -0.0 + logits[:, 9] = 0.0 + + +@pytest.mark.parametrize("ties", [False, True], ids=["peer_rows", "peer_ties"]) +def test_signed_zero_maxima_take_the_lower_index(ties) -> None: + """-0.0 and +0.0 compare equal, so the lower index wins as with torch.argmax (also across ranks).""" + spec = dict(op.CELLS[1].spec, ties=ties) + args = _case(spec, seed_of(f"signed-zero-{ties}"), edit_logits=_signed_zero_maxima) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..fab3b7fd7853 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_slot_scatter entry (the spec sampler's per-slot stores in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_slot_scatter import ( # noqa: E501 + m3_slot_scatter as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" + +STALE = -4 + + +def _case(spec, seed): + """Worker outputs and the store's transposed [slots, width] views of its [width, slots, 1] buffers.""" + g = gen(seed) + n, nslots = spec["n"], spec["slots"] + + def rnd(shape): + return torch.randint(0, VOCAB, shape, generator=g, device="cuda", dtype=torch.int32) + + store_new = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_next = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_lens = torch.full((nslots,), STALE, dtype=torch.int32, device="cuda") + store_draft = torch.full((nslots, spec["store_d"]), STALE, dtype=torch.int32, device="cuda") + slots = torch.randperm(nslots, generator=g, device="cuda")[:n] + lens = torch.randint(1, 5, (n,), generator=g, device="cuda", dtype=torch.int32) + return (slots, rnd((n, spec["src_w"])), store_new.squeeze(-1).T, rnd((n, spec["src_w"])), + store_next.squeeze(-1).T, lens, store_lens, rnd((n, spec["src_d"])), store_draft) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, zero padding included, other slots untouched.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in (("da", args[2]), ("db", args[4]), ("dl", args[6]), ("dd", args[8])): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_repeated_slots() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0] = args[0].clone() + args[0][1] = args[0][0] + with pytest.raises(AssertionError, match="distinct"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py new file mode 100644 index 000000000000..c4ed9c08c2c3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for m3_spec_offsets (RuntimeHooks.spec_offsets_op): the overlap scheduler's two speculative step corrections +in one plain launch, bit for bit the in-place ATen adds, eager and in a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_spec_offsets import ( # noqa: E501 + m3_spec_offsets as op, +) + +assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" + +I32 = torch.iinfo(torch.int32) + + +def _inputs(tokens, seqs, seed, wrap=False): + g = torch.Generator(device="cuda").manual_seed(seed) + lo, hi = (I32.min, I32.max) if wrap else (0, 1 << 20) + pos = torch.randint(lo, hi, (tokens,), generator=g, device="cuda", dtype=torch.int32) + kv = torch.randint(lo, hi, (seqs,), generator=g, device="cuda", dtype=torch.int32) + pos_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (tokens,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + kv_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (seqs,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + if wrap: + pos[:2], pos_off[:2] = I32.max, 1 # overflow wraps, as the ATen add does + kv[:1], kv_off[:1] = I32.min, -1 + return pos, pos_off, kv, kv_off + + +CELLS = [c.spec for c in op.CELLS] + [ + dict(tokens=1, seqs=1), + dict(tokens=300, seqs=75), + dict(tokens=4, seqs=0), +] + + +@pytest.mark.parametrize("spec", CELLS, ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items())) +def test_bitwise_aten_adds(spec) -> None: + for trial in range(3): + pos, pos_off, kv, kv_off = _inputs( + spec["tokens"], spec["seqs"], 97 * trial + spec["tokens"], spec.get("wrap", False) + ) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + assert op(pos, pos_off, kv, kv_off) is True + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"trial {trial}" + + +def test_graph_replays_read_the_live_offsets() -> None: + pos, pos_off, kv, kv_off = _inputs(32, 8, 5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + for step in range(3): + fresh = _inputs(32, 8, 11 + step) + for dst, src in zip((pos, pos_off, kv, kv_off), fresh): + dst.copy_(src) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"step {step}" + + +def test_runner_views_are_accepted() -> None: + """The runner passes position_ids[0, ctx:] and kv_lens_cuda[ctx_reqs:seqs]: 1-D slices of contiguous buffers.""" + position_ids = torch.arange(40, dtype=torch.int32, device="cuda").view(1, 40) + kv_lens = torch.arange(12, dtype=torch.int32, device="cuda") + pos_off = torch.full((40,), 2, dtype=torch.int32, device="cuda") + kv_off = torch.full((12,), 3, dtype=torch.int32, device="cuda") + assert op(position_ids[0, 8:], pos_off[:32], kv_lens[2:10], kv_off[:8]) is True + torch.cuda.synchronize() + assert torch.equal( + position_ids[0], torch.cat([torch.arange(8), torch.arange(8, 40) + 2]).int().cuda() + ) + assert torch.equal( + kv_lens, + torch.cat([torch.arange(2), torch.arange(2, 10) + 3, torch.arange(10, 12)]).int().cuda(), + ) + + +@pytest.mark.parametrize( + "bad", + ["int64", "strided", "length", "2d"], +) +def test_declines_other_layouts_without_writing(bad) -> None: + pos, pos_off, kv, kv_off = _inputs(8, 2, 3) + if bad == "int64": + kv = kv.long() + elif bad == "strided": + pos = torch.zeros(16, dtype=torch.int32, device="cuda")[::2] + elif bad == "length": + pos_off = pos_off[:7] + else: + pos = pos.view(2, 4) + before = (pos.clone(), kv.clone()) + assert op(pos, pos_off, kv, kv_off) is False + torch.cuda.synchronize() + assert torch.equal(pos, before[0]) and torch.equal(kv, before[1]) + + +def test_launches_plainly() -> None: + """It writes step inputs that pre-wait readers take before their grid wait (MSA contract, dispatch + PRODUCER_PREWAIT_STEP_INPUTS): its kernel node carries no programmatic (PDL) serialization.""" + from cuda.bindings import runtime as rt + + pos, pos_off, kv, kv_off = _inputs(32, 8, 7) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + err, _, n = rt.cudaGraphGetNodes(raw, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + flags = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind == rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + flags.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + assert flags == [0], flags diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..1c23a9357876 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_gather entry (the overlap scheduler's step-input gathers in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_step_gather import ( # noqa: E501 + m3_step_gather as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" + +#: Rows of the engine's flat buffers (the max batch), and their stale fill. +MAX_ROWS = 8 +STALE = -4 + + +def _case(spec, seed): + g = gen(seed) + p, t, d, nslots, b0 = spec["p"], spec["t"], spec["d"], spec["slots"], spec["b0"] + i32 = dict(dtype=torch.int32, device="cuda") + new_tok = torch.randint(0, VOCAB, (t, nslots, 1), generator=g, device="cuda", dtype=torch.int32) + if spec.get("two_d"): + new_tok = new_tok.squeeze(-1) + next_draft = torch.randint(0, VOCAB, (nslots, d), generator=g, device="cuda", dtype=torch.int32) + lens = torch.randint(1, t + 1, (nslots,), generator=g, device="cuda", dtype=torch.int32) + slots = torch.randperm(nslots, generator=g, device="cuda")[:p].to(torch.int32) + ppi = slots.repeat_interleave(t) + ids = torch.full((MAX_ROWS * t,), STALE, **i32) + draft = torch.full((MAX_ROWS * d,), STALE, **i32) + pos_off = torch.full((MAX_ROWS * t,), STALE, **i32) + kv_off = torch.full((MAX_ROWS,), STALE, **i32) + return ( + new_tok, + next_draft, + lens, + slots, + ppi, + p, + t, + d, + ids, + spec["ids0"], + draft, + spec["d0"], + pos_off, + kv_off, + b0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in ( + ("ids", args[8]), + ("draft", args[10]), + ("pos_off", args[12]), + ("kv_off", args[13]), + ): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_a_slot_past_the_store() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[3] = args[3].clone() + args[3][0] = 8 + with pytest.raises(AssertionError, match="slots must index the store"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py new file mode 100644 index 000000000000..a4076e46f9d1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py @@ -0,0 +1,228 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_prologue entry (the decode graph's step-input prologue in one launch).""" + +import random + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_step_prologue import ( # noqa: E501 + m3_step_prologue as op, +) +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" +# The reference and the pack tests run this branch's step_inputs (the step epoch, prologue_op): a run that mounts only +# the target's files gets the image's older module, whose graph scatter has no epoch argument. +assert "epoch" in step_inputs._graph_scatter_kernel.arg_names, ( + "tensorrt_llm/_torch/pyexecutor/step_inputs.py is not this branch's (no step epoch): mount the shared pyexecutor " + "files (step_inputs.py, runtime_hooks.py, engine/runners/decoder/runner.py, h2d_cache.py) with the target" +) + +MAX_DESC = StepInputPack.MAX_DESC +WINDOW = StepInputPack.GRAPH_WINDOW +DATA_START = StepInputPack.HEADER_BYTES + StepInputPack.DESC_BYTES * MAX_DESC +ARENA = 1 << 20 + + +def _layout(spec, seed): + """Copies as (dst, src, bytes) arena offsets (a host copy's src: a window offset) and the host copies' bytes.""" + rng = random.Random(seed) + h2d, d2d, data = [], [], [] + used = DATA_START + pos = 0 + for _ in range(spec["h2d"]): + n = rng.choice([1, 3, 4, 7, 16, 24, 32, 48, 64, 5, 12]) + pos += rng.randint(0, 40) # gaps and every alignment + h2d.append((pos, used, n)) + data.append(bytes(rng.getrandbits(8) for _ in range(n))) + pos += n + used += -(-n // 16) * 16 + assert pos < ARENA // 2 and used <= WINDOW + # Device copies write the last quarter; unchained ones read the third. + dpos, spos = 3 * ARENA // 4, ARENA // 2 + for i in range(spec["d2d"]): + if spec.get("chained") and h2d: + dst_h, _, n_h = h2d[i % len(h2d)] + off = rng.randint(0, n_h - 1) + src, n = dst_h + off, n_h - off + else: + n = rng.choice([4, 8, 16, 32, 6, 64]) + spos += rng.randint(0, 24) + src = spos + spos += n + dpos += rng.randint(0, 24) + d2d.append((dpos, src, n)) + dpos += n + assert dpos < ARENA and spos < 3 * ARENA // 4 + return h2d, d2d, data + + +def _window(layout, base: int) -> torch.Tensor: + """The window of ``layout`` for an arena at device address ``base``.""" + h2d, d2d, data = layout + w = torch.zeros(WINDOW, dtype=torch.uint8) + words = [len(h2d), len(d2d)] + for dst, src, n in h2d: + words += [base + dst, src, n] + for dst, src, n in d2d: + words += [base + dst, base + src, n] + w[: 8 * len(words)] = torch.tensor(words, dtype=torch.int64).view(torch.uint8) + for (_, off, n), b in zip(h2d, data): + w[off : off + n] = torch.frombuffer(bytearray(b), dtype=torch.uint8) + return w.cuda() + + +def _arenas(seed): + a = torch.randint(0, 256, (ARENA,), generator=gen(seed), device="cuda", dtype=torch.int32) + a = a.to(torch.uint8) + return a, a.clone() + + +@pytest.mark.parametrize("with_epoch", [True, False], ids=["epoch", "no_epoch"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, with_epoch, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Byte for byte the default prologue's destinations, header clear and epoch.""" + seed = seed_of(cell.why) + layout = _layout(cell.spec, seed) + arena, ref_arena = _arenas(seed) + window, ref_window = _window(layout, arena.data_ptr()), _window(layout, ref_arena.data_ptr()) + epoch = torch.full((1,), 41, dtype=torch.int64, device="cuda") if with_epoch else None + ref_epoch = epoch.clone() if with_epoch else None + certifying(op) + op(window, epoch, MAX_DESC) + op.reference(ref_window, ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena) + # Each window holds its own arena's addresses: past the consumed header, both are left as staged. + for w, a in ((window, arena), (ref_window, ref_arena)): + assert not w[:16].any(), "the header is consumed" + assert torch.equal(w[16:], _window(layout, a.data_ptr())[16:]) + if with_epoch: + assert int(epoch.item()) == int(ref_epoch.item()) == 42 + + +def test_graph_replays_scatter_each_new_window() -> None: + """Captured once on a static window: every replay writes that replay's window and bumps the epoch once.""" + arena, ref_arena = _arenas(7) + window = torch.zeros(WINDOW, dtype=torch.uint8, device="cuda") + epoch = torch.zeros(1, dtype=torch.int64, device="cuda") + ref_epoch = epoch.clone() + op(window, epoch, MAX_DESC) # the empty window: writes nothing + op.reference(window.clone(), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(window, epoch, MAX_DESC) + for step, spec in enumerate( + [dict(h2d=24, d2d=8, chained=True), dict(h2d=0, d2d=0), dict(h2d=40, d2d=4)] + ): + layout = _layout(spec, 100 + step) + window.copy_(_window(layout, arena.data_ptr())) + graph.replay() + op.reference(_window(layout, ref_arena.data_ptr()), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena), step + assert not window[:16].any(), step + assert int(epoch.item()) == int(ref_epoch.item()) == step + 2, step + + +def _pack_case(prologue_op): + pack = StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=True, + in_graph=True, + early_copy=True, + step_epoch=True, + prologue_op=prologue_op, + ) + a, b, c = (torch.zeros(16, dtype=torch.int32, device="cuda") for _ in range(3)) + out = torch.zeros(48, dtype=torch.int32, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _pack_step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + + def host(v): + return torch.tensor(v, dtype=torch.int32).pin_memory() + + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, host(values[0])) + step_inputs.copy_(b, host(values[1])) + if overlap: + step_inputs.copy_(a[:4], host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if not flush: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_pack_with_the_one_launch_prologue_matches_the_default() -> None: + """The same pack steps through the default prologue and the one-launch one: the same outputs and epochs, and the + one-launch prologue is engaged (the default's warm-up compile is not what bumped its epoch).""" + packs = [_pack_case(None), _pack_case(op)] + assert packs[0][0].prologue_op is None and packs[1][0].prologue_op is op + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + outs = [] + for pack, graph, dsts, out in packs: + # A replay with a new pack, one with an eagerly issued pack, and one with nothing to scatter. + _pack_step(pack, graph, dsts, (va, vb), overlap=step == 2, flush=step != 3) + torch.cuda.synchronize() + outs.append((out.tolist(), int(pack.epoch.item()))) + assert outs[0] == outs[1], step + assert outs[1][1] == step + 2, step # one bump at warm-up, then one per replay + + +def _kernel_nodes_launch_plainly(capture) -> list: + """The PDL attributes of every kernel node of a graph ``capture(stream)`` records: the step-input writers must be + plain launches (MSA pre-wait contract: no programmatic serialization or event on them).""" + from cuda.bindings import runtime as rt + + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + capture() + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + # cuda-python changed the signature: (graph[, numNodes]) on older bindings, (graph, nodes, numNodes) on newer ones. + try: + err, _, n = rt.cudaGraphGetNodes(raw) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + except TypeError: + err, _, n = rt.cudaGraphGetNodes(raw, None, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, None, n) + found = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind != rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + continue + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + # Older bindings return no attribute for a kernel launched without one (a plain launch). + found.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + return found + + +def test_the_prologue_launches_plainly() -> None: + window = torch.zeros(WINDOW, dtype=torch.uint8, device="cuda") + epoch = torch.zeros(1, dtype=torch.int64, device="cuda") + op(window, epoch, MAX_DESC) + torch.cuda.synchronize() + flags = _kernel_nodes_launch_plainly(lambda: op(window, epoch, MAX_DESC)) + assert flags == [0], flags diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..e578e1f8e55b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic layer's feed-forward exit on Eagle3 capture layers, in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``). + +Every layer's MoE / MLP output goes unreduced to the fused boundary (all-reduce + residual + the next layer's input +norm, the final norm after the last layer), and a capture layer captures the reduced residual sum that boundary +writes, with no residual of its own. CPU only: the attention, the feed-forward and the boundaries are stand-ins. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) + +LAYERS = 4 + + +class _Spec: + def __init__(self, layers): + self.layers = set(layers) + + def is_layer_capture(self, i): + return i in self.layers + + +@pytest.fixture(params=["prefill", "decode"]) +def layer(request, monkeypatch): + """(core, the target's generic layer bound to a target over ``core``).""" + calls = [] + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + c = SimpleNamespace( + w=w, num_layers=LAYERS, moe_layers=frozenset({2, 3}), rms_eps=1e-6, calls=calls + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, + "_native_allreduce_norm", + lambda core, x, r, norm: (("normed", x, norm), ("updated", x, r)), + ) + + def boundary(self, x, x2, residual, norm_w): + calls.append((x, x2, norm_w)) + return ("next", norm_w), ("sum", x, x2) + + if request.param == "prefill": + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i))) + monkeypatch.setattr(cls, "_swiglu_mlp", lambda self, h, prefix, *a, **k: ("mlp", prefix)) + else: + cls = modeling.DecodeTarget + monkeypatch.setattr( + cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i), False, False) + ) + monkeypatch.setattr( + cls, "_swiglu_mlp", lambda self, h, prefix, *a, **k: (("mlp", prefix), False) + ) + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + monkeypatch.setattr(cls, "_generic_boundary", boundary) + target = cls(c) + return c, target._layer if request.param == "prefill" else target._generic_layer + + +@pytest.mark.parametrize("i", range(LAYERS)) +def test_a_capture_layer_captures_the_boundary_residual_sum(layer, i) -> None: + core, run = layer + captured = [] + out = run( + i, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([i]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + norm = core.w[f"l{i + 1}_in_norm"] if i + 1 < LAYERS else "final" + assert core.calls[-1][2] == norm, ( + "the boundary applies the next layer's norm (the final norm after the last)" + ) + assert captured == [(i, out[1], None)] + if i in core.moe_layers: + assert core.calls[-1][1] == ("routed", i), "the routed part folds into the boundary" + + +def test_a_layer_without_capture_captures_nothing(layer) -> None: + _, run = layer + captured = [] + run( + 1, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([]), + lambda *a: captured.append(a), + ) + assert not captured diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..fd25081cfe12 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,479 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp8 (Rubin, two trays) decode target's fast-path table (``modeling.py``: the row ranges and the selection +functions): its row ranges against the catalog cells that certify them, and the selection for every decode row count. + +Every range's covered shapes are listed in ``ENABLED``; a range not listed covers nothing. + +CPU only: nothing here launches a kernel. +""" + +import inspect +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_draft_step as draft_step_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + triton_sparse_decode as tsd, +) + +ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 2) # every decode row count and the first one past it + + +@pytest.fixture(autouse=True) +def _policy(): + """Every test starts outside require mode with nothing logged, and leaves it so.""" + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_every_covered_shape_is_a_certified_cell(fast_path) -> None: + uncertified = fast_path.keys - fast_path.certified() + assert not uncertified, ( + f"{fast_path.path} covers {sorted(uncertified, key=str)} but " + f"{type(fast_path.entry).__name__}.CELLS certifies none of them: add a cell or narrow the range" + ) + + +def test_the_fast_path_names_are_unique() -> None: + names = [r.path for r in dispatch.FAST_PATHS] + assert len(names) == len(set(names)) + + +def test_the_decode_target_takes_the_gen_workers_batch() -> None: + """Two requests verifying 1 + 3 drafts; a wider step is the prefill target's.""" + assert dispatch.DECODE_MAX_TOKENS == 8 + assert [dispatch.is_decode_step(n) for n in ROWS] == [True] * 8 + [False] + + +# ------------------------------------------------------------------ the cells turned on + + +def test_the_moe_engine_takes_every_decode_step() -> None: + """The MoE lane's m3_moe engine (TP8 width 384, shared expert folded in) at 1..8 rows, unless killed; its sum + pushes into the multicast boundary all-reduce at the engine's rows while that all-reduce is on.""" + expected = set() if dispatch.KILL_MOE_ENGINE else set(range(1, 9)) + assert set(dispatch.MOE_ENGINE.keys) == expected + for rows in ROWS: + want = dispatch.MOE_ENGINE.path if rows in expected else None + assert dispatch.moe_path(rows) == want, rows + push_rows = dispatch.MOE_ENGINE_MAX_TOKENS if dispatch.MC_AR_TP2 else 0 + assert dispatch.MOE_PUSH_MAX_TOKENS == push_rows + assert set(dispatch.MOE_PUSH.keys) == set(range(1, push_rows + 1)) + + +# ------------------------------------------------------------------ on: the sparse layers' block selection + + +def test_the_index_scorer_takes_one_and_two_requests_with_their_schedules() -> None: + assert dispatch.INDEX_SCORE.keys == frozenset({1, 2}) + assert dispatch.index_score_schedule(1) == (1272, True) + assert dispatch.index_score_schedule(2) == (768, True) + # every covered request count has a split (the graph captures it) + assert set(dispatch.INDEX_SCORE.keys) <= set(dispatch.INDEX_SCORE_SPLIT_K) + + +def test_the_long_select_takes_every_key_width_up_to_its_maximum() -> None: + assert dispatch.LONG_SELECT.keys == frozenset(dispatch.LONG_SELECT_KEY_WIDTHS) + assert (dispatch.LONG_SELECT_MIN_BLOCKS, dispatch.LONG_SELECT_MAX_BLOCKS) == (129, 9216) + + +def test_the_384_thread_select_takes_every_key_width_with_its_launch_key() -> None: + assert dispatch.SELECT_LONG_DSL is True + assert dispatch.LONG_SELECT_DSL.keys == frozenset(dispatch.LONG_SELECT_DSL_KEY_WIDTHS) + keys = dispatch.long_select_dsl_keys() + assert [key[1:] for key in keys] == [(False, dispatch.SELECT_LONG_TRIGGER_PHASE)] + + +@pytest.mark.parametrize("num_blocks", [129, 1294, 5888, 9216]) +def test_a_long_select_decline_inside_its_range_is_a_miss(num_blocks) -> None: + """The kernel declining operands its cells certify must not fall back silently under require.""" + assert dispatch.long_select_covers(num_blocks) + # Outside require: logged once, the in-tree selector runs. + dispatch.long_select_fallback(num_blocks, 8) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="long_select"): + dispatch.long_select_fallback(num_blocks, 8) + + +@pytest.mark.parametrize("num_blocks", [1, 16, 128, 9217, 20000]) +def test_rows_outside_the_long_select_range_are_uncovered_even_under_require(num_blocks) -> None: + """Short rows and rows past LONG_SELECT_MAX_BLOCKS take the in-tree selector in every mode (long contexts must not + fail under require).""" + assert not dispatch.long_select_covers(num_blocks) + dispatch.set_require(True) + dispatch.long_select_fallback(num_blocks, 8) + + +def test_taken_logs_without_raising_under_require() -> None: + dispatch.set_require(True) + dispatch.taken(dispatch.INDEX_SCORE.path, 1) + dispatch.taken(dispatch.LONG_SELECT.path, 8) + + +def test_an_index_score_cell_of_another_schedule_certifies_nothing() -> None: + spec = dict(seq_lens=(165521,), qlen=4, heads=1, split_k=512, early=True) + assert tuple(dispatch.INDEX_SCORE.certifies(spec)) == () + assert tuple(dispatch.INDEX_SCORE.certifies(dict(spec, split_k=1272, heads=2))) == () + assert tuple(dispatch.INDEX_SCORE.certifies(dict(spec, split_k=1272))) == (1,) + + +#: Every TP8 per-rank MXFP8 projection (N, K): sparse-layer QKV + index, dense QKV, o_proj, shared gate_up / down, +#: dense gate_up / down. +TP8_PROJECTIONS = ( + (1536, 6144), + (1280, 6144), + (6144, 1024), + (768, 6144), + (6144, 384), + (3072, 6144), + (6144, 1536), +) +_DECODE_ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 1) + +#: Fast paths a lane has turned on, each with the shapes it covers (the rest of the skeleton stays off). +ENABLED = { + # SpecTail: sampled (temperature + top-p) verify of 1-2 requests at draft length 3 (the shard kernel). + "m3_sampled_accept": frozenset({(3, 1), (3, 2)}), + # SpecTail: greedy verify (shard argmax + accept) and the fused draft loop of 1-2 requests at draft length 3, + # and that loop's draft-input rows (1 / 2 later-pass rows, 4 / 8 first-pass rows). + "m3_shard_accept": frozenset({(3, 1), (3, 2)}), + "m3_draft_step": frozenset({(3, 1), (3, 2)}), + "m3_draft_input": frozenset({1, 2, 4, 8}), + # SpecTail: the drafter's BF16 linears on FlashInfer's low-M kernels (TP8 port) at the certified b1 rows. + "m3_draft_lowm": frozenset( + { + ((1280, 12288), 1), + ((1280, 12288), 4), + ((6144, 1024), 1), + ((6144, 1024), 4), + ((4608, 6144), 1), + ((4608, 6144), 4), + ((6144, 2304), 1), + ((6144, 2304), 4), + ((6144, 18432), 1), + } + ), + # SpecTail: the fused first drafter inputs of 1-2 fused-verify requests, the capture norm and the chunk norm at + # their 4 / 8 verify rows (own ranges and kills since the TP2 fast-path audit). + "m3_fused_first_inputs": frozenset({1, 2}), + "m3_capture_norm": frozenset({4, 8}), + "m3_chunk_norm": frozenset({4, 8}), + # MSA: the one-launch long-context sparse decode at the TP8 heads (pre-wait up to its rows, plain above), the + # identity-selection producer, and the sparse layers' NVFP4 QKV + indexer producer at the TP8 heads. + "m3_sparse_decode_lc8": frozenset({"plain", "prewait"}), + "m3_producer": frozenset({"nvfp4_tp8"}), + "sparse_decode_producer": frozenset({"multi_query", "per_row"}), + # MoE: the m3_moe engine (TP8 width 384, shared expert folded in) at 1..8 rows, unless killed. + "m3_moe_engine": frozenset() if dispatch.KILL_MOE_ENGINE else frozenset(range(1, 9)), + # The TF32 tensor-core router of 9..16-row pure generation steps on the generic layer (above the engine's rows). + "m3_router_tc": frozenset(range(9, 17)), + # Indexer: the index scorer at 1-2 requests, the long-row select at every key width (129..9216 blocks). + "m3_index_score": frozenset({1, 2}), + "long_select": frozenset({2, 4, 9}), + "long_select_dsl": frozenset({4, 11, 24}), + # LongAttn: the dense layers' full-context decode of 1-2 requests, 1 + 3 verify rows or one draft-pass row. + "m3_fullctx_decode": frozenset({(4, 1), (1, 1), (4, 2), (1, 2)}), + # GEMM/AR: the decode MXFP8 GEMVs at every TP8 projection and decode row, on an MXFP8 and a bf16 activation. + "m3_mx_gemv": frozenset((nk, m) for nk in TP8_PROJECTIONS for m in _DECODE_ROWS) + if dispatch.MX_GEMV_TP2 + else frozenset(), + "m3_mx_gemv_bf16": frozenset((nk, m) for nk in TP8_PROJECTIONS for m in _DECODE_ROWS) + if dispatch.MX_GEMV_TP2 + else frozenset(), + # GEMM/AR: the multicast boundary all-reduce, the o_proj's push into it, the fused embedding + first RMSNorm. + "mc_allreduce_norm": frozenset(_DECODE_ROWS) if dispatch.MC_AR_TP2 else frozenset(), + "attn_push": frozenset(_DECODE_ROWS) + if (dispatch.MC_AR_TP2 and dispatch.ATTN_PUSH_TP2) + else frozenset(), + "dense_push": frozenset(_DECODE_ROWS) + if (dispatch.MC_AR_TP2 and dispatch.DENSE_PUSH_TP2) + else frozenset(), + # MoE: the engine's push into the multicast boundary all-reduce, at the engine's rows. + "moe_push": frozenset(range(1, dispatch.MOE_PUSH_MAX_TOKENS + 1)), + "embed_firstnorm": frozenset(_DECODE_ROWS) if dispatch.EMBED_FIRSTNORM_TP2 else frozenset(), + # SpecTail: the fused draft loop's all-reduce boundaries (1..8 rows) on the multicast all-reduce (opts 16). + "m3_draft_mc_allreduce_norm": frozenset(range(1, dispatch.DRAFT_MC_AR_MAX_ROWS + 1)), + # LongAttn: the Eagle3 drafter layer's generation passes of 1-2 requests on the same recipes. + "m3_drafter_fullctx": frozenset({(4, 1), (1, 1), (4, 2), (1, 2)}), +} + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_skeleton_no_fast_path_covers_a_shape(fast_path) -> None: + expected = ENABLED.get(fast_path.path, frozenset()) + assert fast_path.keys == expected, f"{fast_path.path} covers {sorted(fast_path.keys, key=str)}" + + +def test_the_sparse_decode_writes_the_mxfp8_handoff() -> None: + """The sparse decode (LC8, the Triton merge) writes the MXFP8 o_proj hand-off; the range keys the target's + SPARSE_DECODE_LC8 / SPARSE_DECODE_PRODUCER cover are the kernel module's.""" + assert tsd.MERGE_MXFP8_HANDOFF is True + assert dispatch.SPARSE_DECODE_LC8.keys is tsd.LC8_KEYS + assert dispatch.SPARSE_DECODE_PRODUCER.keys is tsd.PRODUCER_KEYS + + +def test_sparse_decode_lc8_covers_the_decode_steps_with_the_handoff() -> None: + """Every decode step (at most DECODE_MAX_TOKENS rows) writes the MXFP8 o_proj hand-off; the one-launch decode + stages its pre-wait reads up to m3_attn_lc8.PREWAIT_MAX_ROWS rows and merges once.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_attn_lc8, + ) + + assert dispatch.SPARSE_DECODE_LC8.keys == frozenset({"plain", "prewait"}) + assert tsd.MERGE_MXFP8_MIN_TOKENS <= 1 and dispatch.DECODE_MAX_TOKENS <= tsd.MergeMxfp8.MAX_ROWS + assert ( + m3_attn_lc8.PREWAIT_MAX_ROWS == 8 and m3_attn_lc8.MERGE_ONCE and m3_attn_lc8.ONE_DIV_MERGE + ) + + +def test_the_dense_layers_take_the_fullctx_decode_at_the_gen_workers_shapes() -> None: + """One or two requests, 1 + 3 verify rows or one draft-pass row, each at its own pinned recipe.""" + assert dispatch.FULLCTX_DECODE_ENABLED + assert dispatch.FULLCTX_DECODE.keys == {(4, 1), (1, 1), (4, 2), (1, 2)} + assert set(dispatch.FULLCTX_DECODE_RECIPES) == set(dispatch.FULLCTX_DECODE.keys) + for recipe in dispatch.FULLCTX_DECODE_RECIPES.values(): + assert 2 <= recipe["splits"] <= 128 # the standalone reducer's range + + +def test_the_fullctx_entry_cells_run_at_the_forward_plan_capacity() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention import ( # noqa: E501 + m3_fullctx_decode as entry, + ) + + assert entry.MAX_KV_LEN == dispatch.FULLCTX_DECODE_MAX_KV_LEN + + +def test_the_fullctx_capacity_covers_the_agentx_cache_bound() -> None: + """The forward compares the cache's length bound (max_seq_len plus the speculative rows it reserves) with the plan + capacity; at the AgentX max_seq_len of 1 << 20 with up to 8 draft rows the dense layers must still be taken.""" + assert dispatch.FULLCTX_DECODE_MAX_KV_LEN >= (1 << 20) + 8 + + +def test_the_drafter_takes_the_fullctx_decode_at_the_same_shapes() -> None: + assert dispatch.DRAFTER_FULLCTX_ENABLED + assert dispatch.DRAFTER_FULLCTX.keys == dispatch.FULLCTX_DECODE.keys + + +# ------------------------------------------------------------------ landed: the decode MXFP8 GEMVs (GEMM/AR lane) + + +def test_the_gemv_rows_follow_the_kill_constant() -> None: + expected = set(TP8_PROJECTIONS) if dispatch.MX_GEMV_TP2 else set() + assert set(dispatch.MX_GEMV_ROWS) == expected + assert set(dispatch.MX_GEMV_BF16_MAX_TOKENS) == expected + + +@pytest.mark.parametrize("rows", range(1, dispatch.DECODE_MAX_TOKENS + 1)) +@pytest.mark.parametrize("nk", TP8_PROJECTIONS, ids=str) +def test_every_projection_takes_the_gemv_at_every_decode_row(nk, rows) -> None: + """MXFP8 activation (a producer quantized it) -> the GEMV; bf16 -> the GEMV quantizing in its prologue (staged: + the stock decode-graph GEMMs while MX_GEMV_TP2 is off).""" + on = dispatch.MX_GEMV_TP2 + assert dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) == ( + dispatch.MX_GEMV_PATH if on else dispatch.PREQUANT_GEMM_PATH + ) + assert dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) == ( + dispatch.MX_GEMV_BF16_PATH if on else dispatch.GENERIC_GEMM_PATH + ) + + +@pytest.mark.parametrize("nk", TP8_PROJECTIONS, ids=str) +def test_wider_than_the_decode_step_is_the_generic_gemm(nk) -> None: + rows = dispatch.DECODE_MAX_TOKENS + 1 + assert ( + dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) + == dispatch.PREQUANT_GEMM_PATH + ) + assert ( + dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) + == dispatch.GENERIC_GEMM_PATH + ) + + +def test_an_unknown_projection_is_the_generic_gemm() -> None: + assert ( + dispatch.linear_path(256, 512, 1, prequant=True, low_m_gemm=True) + == dispatch.PREQUANT_GEMM_PATH + ) + assert ( + dispatch.linear_path(256, 512, 1, prequant=False, low_m_gemm=True) + == dispatch.GENERIC_GEMM_PATH + ) + + +def test_taken_logs_once_per_path_and_rows(monkeypatch) -> None: + lines = [] + monkeypatch.setattr(dispatch.logger, "info", lines.append) + for _ in range(3): + dispatch.taken(dispatch.MX_GEMV_PATH, 4, "[6144, 4096] mxfp8 in") + dispatch.taken(dispatch.MX_GEMV_PATH, 8, "[6144, 4096] mxfp8 in") + assert len(lines) == 2 and all("m3_mx_gemv taken at" in line for line in lines) + + +def test_skeleton_the_dense_decode_split_is_the_built_in_one() -> None: + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {} + + +def test_skeleton_speculation_runs_the_in_tree_loop() -> None: + # SpecTail: fused greedy / sampled verify and the fused draft loop at b1/b2 (BF16 head; no coded heads). + assert dispatch.SHARD_ACCEPT_MAX_REQUESTS == 2 + assert dispatch.SAMPLED_VERIFY_MAX_REQUESTS == 2 + assert dispatch.DRAFT_STEP_MAX_REQUESTS == 2 + assert dispatch.FC_EARLY_MAX_ROWS == 0 + assert not draft_step_kernel.HEAD_CODED + + +def test_a_miss_raises_only_under_require() -> None: + dispatch.missed("some_path", 1, "a precondition") + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="some_path"): + dispatch.missed("some_path", 1, "a precondition") + + +def test_skeleton_first_drafter_inputs_are_in_tree() -> None: + # SpecTail turned the sampled verify on (SAMPLED_VERIFY_MAX_REQUESTS 2, test_skeleton_speculation_...). + assert dispatch.FIRST_DRAFTER_INPUTS_MAX_REQUESTS == 0 + + +@pytest.mark.parametrize( + "rows, capture_normed, taken", + [ + (4, False, True), + (8, False, True), + (8, True, True), + (12, False, False), + (12, True, False), + (16, True, False), + ], +) +def test_the_fused_first_inputs_take_only_their_covered_rows(rows, capture_normed, taken) -> None: + """The fused first drafter inputs (and the fc fork) take a fused-verify step only at the requests + FUSED_FIRST_INPUTS covers and, unless the captures applied the fc_norm, the rows CHUNK_NORM covers; anything else + runs the in-tree first drafter inputs.""" + + worker = object.__new__(dispatch.M3Eagle3Worker) + worker._m3_capture_norm = (None, 1e-6) if capture_normed else None + meta = SimpleNamespace(runtime_draft_len=dispatch.SHARD_ACCEPT_DRAFT_LEN) + assert worker._m3_first_inputs_covered(rows, meta) is taken + + +# ------------------------------------------------------- landed: the multicast boundary all-reduce (GEMM/AR lane) + + +def test_both_boundaries_take_the_multicast_all_reduce() -> None: + """One 8-token bucket for the whole decode step; the MoE boundary publishes the next QKV's MXFP8 rows (killed: + every boundary on the built-in all-reduce while MC_AR_TP2 is off).""" + on = dispatch.MC_AR_TP2 + assert tuple(p.max_tokens for p in dispatch.DECODE_PLANS) == (dispatch.DECODE_MAX_TOKENS,) + for plan in dispatch.DECODE_PLANS: + assert (plan.attn_mc_ar, plan.moe_mc_ar, plan.qkv_prequant) == (on, on, on), plan + assert dispatch.ATTN_PUSH_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if on and dispatch.ATTN_PUSH_TP2 else 0 + ) + assert dispatch.DENSE_PUSH_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if on and dispatch.DENSE_PUSH_TP2 else 0 + ) + for rows in _DECODE_ROWS: + assert dispatch.plan_for(rows) is dispatch.DECODE_PLANS[0] + # The MoE pushes only from its engine, at the engine's rows. + assert dispatch.MOE_PUSH_MAX_TOKENS == (dispatch.MOE_ENGINE_MAX_TOKENS if on else 0) + + +def test_the_generic_boundaries_take_the_multicast_all_reduce() -> None: + """Prefill, mixed steps and the Eagle3 capture layers: the same certified rows, alone and folded.""" + rows = dispatch.DECODE_MAX_TOKENS if dispatch.MC_AR_TP2 else 0 + assert dispatch.GENERIC_MC_AR_MAX_TOKENS == rows + assert dispatch.GENERIC_MC_AR_FOLDED_MAX_TOKENS == rows + + +# ------------------------------------------------------------------ landed: the fused embedding + first RMSNorm + + +def test_the_fused_embedding_rows_follow_the_kill_constant() -> None: + assert dispatch.EMBED_FIRSTNORM_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if dispatch.EMBED_FIRSTNORM_TP2 else 0 + ) + + +# ------------------------------------------------------- landed: MX-TAC (pinned decode-graph MXFP8 GEMM tactic) + + +def test_the_decode_mxfp8_decision_is_pinned_to_cute_dsl() -> None: + from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, + ) + + assert dispatch.MXFP8_DECODE_QUANT_TACTIC == MXFP8QuantizeRunner.CUTE_DSL + assert dispatch.MXFP8_DECODE_GEMM_TACTIC == FlashInferMXFP8GemmRunner.CUTE_DSL + + +@pytest.mark.parametrize("target", ["DecodeTarget", "PrefillTarget"]) +def test_the_moe_forward_reports_the_engine_taken(target) -> None: + """The engagement evidence comes from the forward (not the engine build's warm call).""" + + src = inspect.getsource(getattr(dispatch, target)._moe) + assert src.index("taken(MOE_ENGINE.path") < src.index("m3_moe_engine(h.contiguous()") + + +# ------------------------------------------------------------------ the drafter's low-M linears (TP8 port) + + +def test_draft_lowm_covers_exactly_the_certified_rows_of_its_tactic_table() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm import ( # noqa: E501 + m3_draft_lowm as entry, + ) + + assert set(entry.TACTICS) == set(entry.ROWS) == set(dispatch.DRAFT_LOWM_ROWS) + assert dispatch.DRAFT_LOWM.keys == frozenset( + (nk, m) + for nk, rows in entry.ROWS.items() + for m in rows + if m <= dispatch.DRAFT_LOWM_MAX_ROWS + ) + for nk, table in entry.TACTICS.items(): + rows = [row[0] for row in table] + assert rows == sorted(rows) and len(set(rows)) == len(rows), nk + for m in entry.ROWS[nk]: + assert entry.tactic_for(*nk, m) is not None, (nk, m) + + +def test_prebuild_draft_head_is_idempotent_and_keeps_replaced_copies(monkeypatch) -> None: + """A repeated cache_derived_state with an unchanged head weight keeps the built (placed) copy; a changed weight + rebuilds it and keeps the old copy alive, so a CUDA graph captured earlier never reads unmapped memory.""" + + import torch + + built = [] + + def fake_prepare(lm_head, tp_rank): + head = SimpleNamespace(weight=lm_head.weight) + built.append(head) + return head + + monkeypatch.setattr(draft_step_kernel, "prepare_head", fake_prepare) + worker = object.__new__(dispatch.M3Eagle3Worker) + worker._m3_draft_head = None + worker._m3_draft_head_version = None + worker._m3_draft_head_graveyard = [] + worker.mapping = SimpleNamespace(tp_rank=0) + monkeypatch.setattr(worker, "_m3_drafter_supported", lambda dm: True, raising=False) + model = SimpleNamespace(lm_head=SimpleNamespace(weight=torch.zeros(4, 4))) + + assert worker.prebuild_draft_head(model) and len(built) == 1 + first = worker._m3_draft_head + assert worker.prebuild_draft_head(model) and len(built) == 1 and worker._m3_draft_head is first + model.lm_head.weight.add_(1.0) # new values: version bump + assert worker.prebuild_draft_head(model) and len(built) == 2 + assert worker._m3_draft_head is built[1] and worker._m3_draft_head_graveyard == [first] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..9414fae5c97f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode robustness without a GPU kernel: the multicast all-reduce epilogue options keep the cluster-barrier bit, +an engine build that fails on one TP rank drops that engine on every rank (two ranks, and the tp8 target's eight, +simulated with threads), and the Eagle3 fc fork's output is recorded on the stream that reads it.""" + +import threading +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = modeling.require_mode() + yield + modeling.set_require(require) + modeling.reset_logged() + + +@pytest.mark.parametrize("finalize", [False, True]) +@pytest.mark.parametrize("bucket", [8, 16, 32, 48]) +def test_ar_opts_keep_cluster_barrier_bit(monkeypatch, bucket, finalize): + # Tables without bit 2 and with buckets missing: every lookup, fallback included, keeps the barrier. + monkeypatch.setattr(mc, "AR_OPTS_BY_M", (0, 1, 4, 5, 8, 9, 12, 13)) + monkeypatch.setattr(mc, "AR_OPTS_BY_BUCKET", {}) + monkeypatch.setattr(mc, "AR_FIN_OPTS_BY_BUCKET", {}) + for m in range(1, min(bucket, 8) + 1): + opts = mc._ar_opts(m, bucket, finalize=finalize) + assert opts & mc._OPT_CLUSTER_BARRIER and 0 <= opts <= 15, (m, bucket, finalize, opts) + + +def test_ar_opts_shipped_tables_unchanged(): + for m in range(1, 9): + assert mc._ar_opts(m, 8) == mc.AR_OPTS_BY_M[m - 1] + for b, o in mc.AR_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b) == o + for b, o in mc.AR_FIN_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b, finalize=True) == o + + +class _FakeComm: + """An MPI-like TP communicator for ``world`` threads: allreduce sums the ranks' ints.""" + + def __init__(self, world): + self.world = world + self.barrier = threading.Barrier(world) + self.lock = threading.Lock() + self.vals = [] + self.calls = 0 + + def allreduce(self, v): + with self.lock: + self.vals.append(v) + self.calls += 1 + self.barrier.wait() + total = sum(self.vals) + self.barrier.wait() + with self.lock: + self.vals.clear() + self.barrier.wait() + return total + + def Get_size(self): # noqa: N802 (MPI name) + return self.world + + +def _mapping(tp_size=2, adp=False): + return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) + + +def _run_ranks(fns): + out, errs = [None] * len(fns), [] + + def run(i): + try: + out[i] = fns[i]() + except BaseException as e: # noqa: BLE001 + errs.append(e) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errs, errs + return out + + +def _patch_comm(monkeypatch, world: int) -> _FakeComm: + from tensorrt_llm._torch.distributed import ops + + comm = _FakeComm(world) + monkeypatch.setattr(mc, "_tp_comm", lambda mapping: comm) + monkeypatch.setattr(ops, "mpi_disabled", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return comm + + +@pytest.fixture +def two_rank_comm(monkeypatch): + return _patch_comm(monkeypatch, 2) + + +def test_tp_agree_single_rank_and_adp_take_no_collective(monkeypatch): + monkeypatch.setattr(mc, "_tp_comm", mock.Mock(side_effect=AssertionError("no collective"))) + assert mc.tp_agree(_mapping(1), True) is True + assert mc.tp_agree(_mapping(1), False) is False + assert mc.tp_agree(_mapping(8, adp=True), False) is False + + +def test_tp_agree_refuses_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError): + mc.tp_agree(_mapping(2), True) + + +@pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) +def test_tp_agree_built_two_ranks(two_rank_comm, oks): + built = [object() if ok else None for ok in oks] + out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + assert out == (built if all(oks) else [None, None]) + + +def _core(fail: bool, world: int = 2) -> SimpleNamespace: + """The engine-building half of the core (``MiniMaxM3Tp8ModelingV2``) on one of ``world`` simulated TP ranks, + its engine builds stubbed (and failing where ``fail``).""" + core = SimpleNamespace(mapping=_mapping(world), moe_engines={}) + cls = modeling.MiniMaxM3Tp8ModelingV2 + + def build(i, num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + core._build_engine = build + for name in ("_prebuild_engines", "_try_build", "engine", "_lookup"): + setattr(core, name, getattr(cls, name).__get__(core)) + return core + + +@pytest.mark.parametrize( + "world,failing", + [(2, ()), (2, (1,)), (8, ()), (8, (5,)), (8, (0, 7))], + ids=["w2_ok", "w2_rank1_fails", "w8_ok", "w8_rank5_fails", "w8_ranks0_7_fail"], +) +def test_engine_build_agreed_across_ranks(monkeypatch, world, failing): + comm = _patch_comm(monkeypatch, world) + cores = [_core(rank in failing, world) for rank in range(world)] + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + keys = sorted(cores[0].moe_engines[3]) + assert keys and comm.calls == world * len(keys) + assert not any(hasattr(c, "moe16_engines") for c in cores) + for c in cores[1:]: + assert c.moe_engines[3].keys() == cores[0].moe_engines[3].keys() + for key in keys: + for c in cores: + assert (c.moe_engines[3][key] is None) == bool(failing), key + # Every rank then takes the same path in the forward: the engine, or (after a miss) the generic MoE. + modeling.set_require(False) + for n in sorted(modeling.MOE_ENGINE.keys): + got = [c.engine(3, n, shared=True) is None for c in cores] + assert len(set(got)) == 1, (n, got) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + M3Eagle3Worker = modeling.M3Eagle3Worker + main, fork = mock.Mock(name="main"), mock.Mock(name="fork") + hidden = mock.Mock(name="hidden", is_cuda=True) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *a: main) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + worker = SimpleNamespace(_m3_fc_hidden=hidden, _m3_fc_flagged=False) + M3Eagle3Worker._m3_fc_join(worker, fork) + main.wait_stream.assert_called_once_with(fork) + if capturing: + hidden.record_stream.assert_not_called() + else: + hidden.record_stream.assert_called_once_with(main) + worker._m3_fc_hidden = None + M3Eagle3Worker._m3_fc_join(worker, fork) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index 1f0a0ea6c085..0701abb94239 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -554,16 +554,84 @@ def test_routes_on_the_kv_cache_algorithm_the_engine_runs( # experts split by expert rather than by tensor: another weight layout, no target (dict(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8), None), (dict(world_size=8, tp_size=8, moe_tp_size=8, enable_attention_dp=True), None), - (dict(world_size=2, tp_size=2, moe_tp_size=2), None), + (dict(world_size=2, tp_size=2, moe_tp_size=2, moe_ep_size=1), "tp2"), + (dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2), None), ], ) def test_minimax_m3_parallel_names_the_target_topology(mapping_kwargs, parallel): - """MiniMax-M3 has one target per tensor-parallel width: tp4 and tp8 load different per-rank shapes.""" + """MiniMax-M3 has one target per tensor-parallel width: tp2, tp4 and tp8 load different per-rank shapes.""" from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing assert routing._parallel(Mapping(**mapping_kwargs)) == parallel +_M3_TARGETS = "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3" +_M3_TP2 = dict(world_size=2, tp_size=2, moe_tp_size=2, moe_ep_size=1) +_M3_TP4 = dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1) +_M3_TP8 = dict(world_size=8, tp_size=8, moe_tp_size=8, moe_ep_size=1) + + +@pytest.mark.parametrize( + "capability, mapping_kwargs, target", + [ + ((10, 3), _M3_TP4, "minimax_m3_nvfp4__sm_103__tp4.ModelingV2MiniMaxM3Nvfp4Sm103Tp4"), + ((10, 3), _M3_TP8, "minimax_m3_nvfp4__sm_103__tp8.ModelingV2MiniMaxM3Nvfp4Sm103Tp8"), + ((10, 7), _M3_TP2, "minimax_m3_nvfp4__sm_107__tp2.ModelingV2MiniMaxM3Nvfp4Sm107Tp2"), + ((10, 7), _M3_TP4, "minimax_m3_nvfp4__sm_107__tp4.ModelingV2MiniMaxM3Nvfp4Sm107Tp4"), + ((10, 7), _M3_TP8, "minimax_m3_nvfp4__sm_107__tp8.ModelingV2MiniMaxM3Nvfp4Sm107Tp8"), + # no GB300 tp2 target + ((10, 3), _M3_TP2, None), + ], +) +def test_minimax_m3_route_picks_the_target_by_architecture_and_width( + monkeypatch, capability, mapping_kwargs, target +): + """Each (SM, tensor-parallel width) has its own target directory. The target modules are stubbed: this checks the + decision, not that the modeling code imports on this host.""" + import sys + import types + + from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + from tensorrt_llm.quantization.mode import QuantAlgo + + for directory, name in ( + ("minimax_m3_nvfp4__sm_103__tp4", "ModelingV2MiniMaxM3Nvfp4Sm103Tp4"), + ("minimax_m3_nvfp4__sm_103__tp8", "ModelingV2MiniMaxM3Nvfp4Sm103Tp8"), + ("minimax_m3_nvfp4__sm_107__tp2", "ModelingV2MiniMaxM3Nvfp4Sm107Tp2"), + ("minimax_m3_nvfp4__sm_107__tp4", "ModelingV2MiniMaxM3Nvfp4Sm107Tp4"), + ("minimax_m3_nvfp4__sm_107__tp8", "ModelingV2MiniMaxM3Nvfp4Sm107Tp8"), + ): + stub = types.ModuleType(f"{_M3_TARGETS}.{directory}.modeling") + setattr(stub, name, type(name, (), {"__module__": stub.__name__})) + monkeypatch.setitem(sys.modules, stub.__name__, stub) + monkeypatch.setattr(routing, "_quant_algo", lambda _config: QuantAlgo.MIXED_PRECISION) + + text_config = dict( + num_hidden_layers=60, + hidden_size=6144, + num_local_experts=128, + intermediate_size=3072, + num_attention_heads=64, + num_key_value_heads=4, + ) + pretrained = types.SimpleNamespace(text_config=text_config) + mapping = Mapping(**mapping_kwargs) + runtime = _m3_runtime() + if target is None: + with pytest.raises(NoTarget, match="sm107 \\(Rubin\\) with 2, 4 or 8 ranks"): + routing.route( + pretrained, mapping, capability, None, runtime.kv_cache_quant_algo, runtime=runtime + ) + return + resolved = routing.route( + pretrained, mapping, capability, None, runtime.kv_cache_quant_algo, runtime=runtime + ) + directory, name = target.split(".") + assert resolved.__module__ == f"{_M3_TARGETS}.{directory}.modeling" + assert resolved.__name__ == name + + # 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, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_one_model_sampling_rejection.py b/tests/unittest/_torch/speculative/hw_agnostic/test_one_model_sampling_rejection.py index a53c2406b55e..c318ae4b7b20 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_one_model_sampling_rejection.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_one_model_sampling_rejection.py @@ -26,6 +26,7 @@ from tensorrt_llm._torch.speculative.spec_sampler_base import ( UNSUPPORTED_BAD_WORDS_MSG, UNSUPPORTED_EMBEDDING_BIAS_MSG, + UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG, UNSUPPORTED_MIN_LENGTH_MSG, UNSUPPORTED_MIN_P_MSG, UNSUPPORTED_NO_REPEAT_NGRAM_MSG, @@ -45,6 +46,11 @@ ("bad", dict(bad="nope"), UNSUPPORTED_BAD_WORDS_MSG), ("no_repeat_ngram_size", dict(no_repeat_ngram_size=3), UNSUPPORTED_NO_REPEAT_NGRAM_MSG), ("top_p_decay", dict(top_p_decay=0.5), UNSUPPORTED_TOP_P_DECAY_MSG), + ( + "logits_processor", + dict(logits_processor=lambda req_id, logits, ids, stream, client_id: None), + UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG, + ), ] # Values a frontend forwards explicitly at their neutral setting. These must not @@ -162,6 +168,11 @@ def test_check_arguments_admits_drafter_modes(): ("min_length", dict(py_min_length=16), UNSUPPORTED_MIN_LENGTH_MSG), ("bad_words", dict(py_bad_words=[[7]]), UNSUPPORTED_BAD_WORDS_MSG), ("no_repeat_ngram_size", dict(py_no_repeat_ngram_size=3), UNSUPPORTED_NO_REPEAT_NGRAM_MSG), + ( + "logits_post_processor", + dict(py_logits_post_processors=[lambda req_id, logits, ids, stream, client_id: None]), + UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG, + ), ] @@ -272,3 +283,21 @@ def test_executor_side_check_still_rejects_top_p_decay(): def test_executor_side_check_admits_a_neutral_request(): """The stub covers every field the path reads; a neutral request passes.""" _validate_on_executor(_executor_request()) + + +@pytest.mark.cpu_only +def test_executor_side_check_admits_an_empty_post_processor_list(): + """An empty list is what a request without callbacks may carry: not rejected.""" + request = _executor_request() + request.py_logits_post_processors = [] + _validate_on_executor(request) + + +@pytest.mark.cpu_only +def test_executor_side_check_rejects_a_post_processor_on_the_full_path(): + """Through the whole validate_request, as the executor admits a request.""" + request = _executor_request() + request.py_logits_post_processors = [lambda req_id, logits, ids, stream, client_id: None] + with pytest.raises(ValueError) as excinfo: + _validate_on_executor(request) + assert str(excinfo.value) == UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py index 0249d62b4efa..9ae828257f9e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -604,7 +604,7 @@ def copy_to_host(tensor): transfers.append(("copy", tuple(tensor.shape))) return tensor.clone() - def record_event(): + def record_event(side_stream_event=None): transfers.append(("record", None)) return types.SimpleNamespace(synchronize=lambda: transfers.append(("sync", None))) diff --git a/tests/unittest/_torch/speculative/test_rng_upload_pack.py b/tests/unittest/_torch/speculative/test_rng_upload_pack.py new file mode 100644 index 000000000000..1e740855ecb5 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_rng_upload_pack.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Philox seed/offset uploads of a sampled one-model step, staged into the step-input pack.""" + +import types +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack +from tensorrt_llm._torch.speculative.interface import SpecMetadata + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.skipif( + not (torch.cuda.is_available() and hasattr(torch.ops.trtllm, "step_inputs_scatter")), + reason="torch.ops.trtllm.step_inputs_scatter is not registered", + ), +] + +MAX_DRAFT_LEN = 3 +ROWS_PER_REQUEST = MAX_DRAFT_LEN + 1 + + +def _meta(pack_rng_uploads: bool) -> SpecMetadata: + meta = SpecMetadata( + max_num_requests=4, max_draft_len=MAX_DRAFT_LEN, max_total_draft_tokens=MAX_DRAFT_LEN + ) + meta.pack_rng_uploads = pack_rng_uploads + return meta + + +def _request(slot: int, seed: Optional[int]) -> types.SimpleNamespace: + return types.SimpleNamespace( + py_seq_slot=slot, + py_decoding_iter=0, + py_request_id=slot, + sampling_config=types.SimpleNamespace(seed=seed), + ) + + +def _populate(meta: SpecMetadata) -> None: + requests = [_request(2, seed=1234), _request(0, seed=None)] + normalized = [(0.7, 0, 0.9, 0.0, ROWS_PER_REQUEST) for _ in requests] + meta._populate_request_rng_state(requests, normalized) + + +def _rng_state(meta: SpecMetadata) -> dict[str, torch.Tensor]: + torch.cuda.synchronize() + return { + name: getattr(meta, name).clone() + for name in ("request_seeds", "request_offsets", "seeds", "offsets") + } + + +def test_packed_uploads_land_at_the_flush_with_the_plain_values(): + plain = _meta(pack_rng_uploads=False) + _populate(plain) + expected = _rng_state(plain) + + packed = _meta(pack_rng_uploads=True) + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + with step_inputs.active(pack): + step_inputs.begin() + _populate(packed) + before_flush = _rng_state(packed) + step_inputs.flush() + after_flush = _rng_state(packed) + + for name, tensor in expected.items(): + assert torch.equal(after_flush[name], tensor), name + # Staged, not copied: the buffers still hold their zero init until the flush. + assert not torch.equal(before_flush["seeds"], expected["seeds"]) + + +def test_without_an_open_pack_the_uploads_are_plain_copies(): + plain = _meta(pack_rng_uploads=False) + _populate(plain) + expected = _rng_state(plain) + + packed = _meta(pack_rng_uploads=True) + _populate(packed) + + for name, tensor in _rng_state(packed).items(): + assert torch.equal(tensor, expected[name]), name + + +def test_packing_is_off_by_default(): + assert not SpecMetadata( + max_num_requests=1, max_draft_len=MAX_DRAFT_LEN, max_total_draft_tokens=MAX_DRAFT_LEN + ).pack_rng_uploads diff --git a/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py index 1ea8cdaa58f4..df9b16704655 100644 --- a/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py +++ b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py @@ -252,3 +252,119 @@ def test_a_fold_that_declines_leaves_the_sampler_store_path(draft_width): assert fold.rows_stored_calls == [(0, SEQ_SLOTS)] assert len(fold.scatter_calls) == 1 _assert_store_equals(sampler, _reference_store(before, outputs)) + + +def _host_outputs(state) -> dict[str, torch.Tensor]: + state.sampler_event.synchronize() + return dict( + new_tokens=state.host.new_tokens.clone(), + new_tokens_lens=state.host.new_tokens_lens.clone(), + next_draft_tokens=state.host.next_draft_tokens.clone(), + ) + + +def test_store_copies_stay_on_the_main_stream_by_default(): + sampler = _sampler() + scheduled, outputs = _step(MAX_DRAFT_LEN) + + state = sampler.sample_async(scheduled, outputs, []) + + assert not sampler.d2h_side_stream_copies + assert state.sampler_event.side_stream_event is None + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_side_stream_store_copies_match_the_main_stream_copies(draft_width): + scheduled, outputs = _step(draft_width) + on_main = _host_outputs(_sampler().sample_async(scheduled, outputs, [])) + sampler = _sampler() + sampler.d2h_side_stream_copies = True + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is not None + on_side = _host_outputs(state) + for name, tensor in on_main.items(): + assert torch.equal(on_side[name], tensor), name + + +def test_a_fold_that_stored_the_rows_keeps_the_copies_on_the_main_stream(): + sampler = _sampler() + sampler.bind_spec_step_fold(_fold(rows_stored=True)) + sampler.d2h_side_stream_copies = True + scheduled, outputs = _step(MAX_DRAFT_LEN) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is None + + +def test_a_fold_that_awaits_the_store_copies_moves_them_to_the_side_stream(): + """A fold with store_copies_done: the copies go on the side stream although its forward stored the rows, and + the event is recorded after them.""" + scheduled, outputs = _step(MAX_DRAFT_LEN) + fold = _fold(rows_stored=True) + fold.store_copies_done = torch.cuda.Event(external=True) + fold.store_copies_done.record() + sampler = _sampler() + sampler.bind_spec_step_fold(fold) + sampler.d2h_side_stream_copies = True + expected = _store_snapshot(sampler) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is not None + state.sampler_event.synchronize() + fold.store_copies_done.synchronize() + host = _host_outputs(state) + for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens"): + assert torch.equal(host[name], expected[name].cpu()), name + + +def test_a_fold_that_awaits_the_store_copies_keeps_them_on_the_main_stream_when_side_copies_are_off(): + scheduled, outputs = _step(MAX_DRAFT_LEN) + fold = _fold(rows_stored=True) + fold.store_copies_done = torch.cuda.Event(external=True) + sampler = _sampler() + sampler.bind_spec_step_fold(fold) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is None + + +def test_static_slot_ids_are_off_by_default(): + assert not _sampler().static_slot_ids + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_static_slot_ids_store_the_same_rows(draft_width): + scheduled, outputs = _step(draft_width) + plain = _sampler() + expected = _reference_store(_store_snapshot(plain), outputs) + sampler = _sampler() + sampler.static_slot_ids = True + + sampler.sample_async(scheduled, outputs, []) + sampler.sample_async(scheduled, outputs, []) # second step reuses the uploaded ids + + _assert_store_equals(sampler, expected) + + +def test_static_slot_ids_follow_a_changed_batch(): + sampler = _sampler() + sampler.static_slot_ids = True + + first_view = sampler._device_slot_ids([4, 1]) + first = first_view.clone() + # Unchanged ids reuse the sampler's buffer without a new upload. The returned tensor is a view of that buffer, + # which the next call rewrites, so snapshot it before changing the batch. + reused_view = sampler._device_slot_ids([4, 1]) + assert reused_view.data_ptr() == first_view.data_ptr() + reused = reused_view.clone() + changed = sampler._device_slot_ids([2, 5, 0]).clone() + + torch.cuda.synchronize() + assert first.tolist() == [4, 1] + assert reused.tolist() == [4, 1] + assert changed.tolist() == [2, 5, 0] diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 0ed31615b378..9fca7902d91e 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -15,6 +15,7 @@ validate_and_set_kv_cache_quant, validate_encoder_decoder_kv_cache_config, ) +from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -189,18 +190,32 @@ def test_validate_and_set_kv_cache_quant_auto_keeps_quant_config_dict(): assert layer_quant_config.kv_cache_quant_algo is None +@pytest.mark.parametrize("implementation", [None, "triton", "msa"]) +@pytest.mark.parametrize("requested_dtype", ["nvfp4", "auto"]) def test_validate_and_set_kv_cache_quant_sm107_downgrade_propagates_to_quant_config_dict( monkeypatch: pytest.MonkeyPatch, + implementation: str | None, + requested_dtype: str, ) -> None: - """The SM107 NVFP4->FP8 downgrade must update per-layer QuantConfigs too, - otherwise a mixed-precision checkpoint keeps NVFP4 in quant_config_dict - while the KV pool (sized from the global config) uses FP8.""" + """The SM107 NVFP4->FP8 decision must reach per-layer QuantConfigs too, + otherwise a mixed-precision checkpoint and the KV pool (sized from the + global config) disagree. MiniMax-M3 MSA keeps NVFP4, which its hybrid + cache serves on SM107; every other configuration falls back to FP8.""" _mock_sm107(monkeypatch) model_config = _make_mixed_precision_model_config() - validate_and_set_kv_cache_quant(model_config, "nvfp4") - assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + if implementation is not None: + model_config.sparse_attention_config = MiniMaxM3SparseAttentionConfig( + implementation=implementation + ) + if requested_dtype == "auto": + model_config.quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4 + for layer_quant_config in model_config.quant_config_dict.values(): + layer_quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4 + validate_and_set_kv_cache_quant(model_config, requested_dtype) + expected = QuantAlgo.NVFP4 if implementation == "msa" else QuantAlgo.FP8 + assert model_config.quant_config.kv_cache_quant_algo == expected for layer_quant_config in model_config.quant_config_dict.values(): - assert layer_quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert layer_quant_config.kv_cache_quant_algo == expected def _write_safetensors_header(checkpoint_dir, tensor_dtype, tensor_shape): diff --git a/tests/unittest/_torch/thop/parallel/test_mxfp8_cute_pinned_tactics.py b/tests/unittest/_torch/thop/parallel/test_mxfp8_cute_pinned_tactics.py new file mode 100644 index 000000000000..f9c38516d2fb --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_mxfp8_cute_pinned_tactics.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The pinned FlashInfer CuTe DSL MXFP8 GEMM tactics (custom_ops/mxfp8_cute_pinned_tactics): every entry is +non-split-K, each +pinned call equals FlashInfer's default (heuristic) CuTe DSL path bit for bit, the runner takes the pinned path above 32 +rows only, and the kill switch restores the default path.""" + +import hashlib + +import pytest +import torch + +from tensorrt_llm._torch.custom_ops import mxfp8_cute_pinned_tactics as pin +from tensorrt_llm._utils import get_sm_version + +sm107 = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 107, reason="SM107 (Rubin) only" +) + + +@pytest.fixture(autouse=True) +def _armed(monkeypatch): + """The table as an opted-in model's executor arms it (RuntimeHooks.mxfp8_pinned_tactics).""" + monkeypatch.setattr(pin, "_ARMED", [True]) + + +def _table_digest() -> str: + return hashlib.md5(repr(sorted(pin.TACTICS.items())).encode()).hexdigest()[:8] + + +def test_table_is_non_split_k_and_pinned(): + assert pin.TACTICS, "empty table" + for (n, k, m), (tile, cluster, swap, prefetch, split_k) in pin.TACTICS.items(): + assert split_k == 1, ((n, k, m), "a split-K tactic changes the K order (not bitwise)") + assert m >= 64 and m & (m - 1) == 0 and m <= pin.MAX_BUCKET, (n, k, m) + assert ( + len(tile) == 2 + and len(cluster) == 2 + and isinstance(swap, bool) + and isinstance(prefetch, bool) + ) + # Drift guard: a re-tuned table is a new measured artifact (bitwise + timing), not an edit in place. + assert _table_digest() == "d43d4a0d", _table_digest() + + +def test_bucket_and_row_gate(): + assert ( + pin._bucket(33) == 64 + and pin._bucket(64) == 64 + and pin._bucket(65) == 128 + and pin._bucket(184) == 256 + ) + assert pin.tactic_for(32, 4992, 6144) is None # 1..32 rows keep FlashInfer's own low-M path + + +def _operands(m, n, k, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, device="cuda", generator=g) * 0.5).to(torch.bfloat16) + w = (torch.randn(n, k, device="cuda", generator=g) * 0.05).to(torch.bfloat16) + q, s = torch.ops.trtllm.mxfp8_quantize(x, True) + wq, ws = torch.ops.trtllm.mxfp8_quantize(w, True) + return q, s, wq, ws + + +def _default(q, s, wq, ws): + from flashinfer import autotune, mm_mxfp8 + + with autotune(tune_mode=False, skip_ops="mxfp8_gemm"): + return mm_mxfp8( + q, wq.t(), s, ws, out_dtype=torch.bfloat16, use_8x4_sf_layout=False, backend="cute-dsl" + ) + + +@sm107 +@pytest.mark.parametrize("key", sorted(pin.TACTICS), ids=lambda k: f"N{k[0]}_K{k[1]}_M{k[2]}") +def test_pinned_equals_default_bitwise(key): + n, k, bucket = key + for m in sorted({bucket // 2 + 1 if bucket > 64 else 33, bucket}): # both ends of the bucket + q, s, wq, ws = _operands(m, n, k, seed=hash(key) % 1000 + m) + assert pin.tactic_for(m, n, k) == pin.TACTICS[key] + out = pin.run(q, s, wq, ws, torch.bfloat16) + assert out is not None, "the pinned path declined a table entry" + assert torch.equal(out, _default(q, s, wq, ws)), (key, m) + + +@sm107 +def test_runner_takes_pinned_above_32_rows_and_kill(monkeypatch): + from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + calls = [] + real = pin.run + monkeypatch.setattr(pin, "run", lambda *a, **kw: calls.append(a[0].shape[0]) or real(*a, **kw)) + for m in (16, 80): + q, s, wq, ws = _operands(m, 4992, 6144, seed=m) + out = runner.forward([q, s, wq, ws], tactic=FlashInferMXFP8GemmRunner.CUTE_DSL) + assert torch.equal(out, _default(q, s, wq, ws)), m + assert calls == [16, 80] + monkeypatch.setattr(pin, "ENABLED", False) + assert pin.tactic_for(80, 4992, 6144) is None From 8b49063871a9b8719ce43f6db577f476b858d5e7 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:31:08 -0700 Subject: [PATCH 12/29] [None][test] Port the kv_lens_update_once test to DecoderInputPreparer The main merge moved DecoderRunner._preprocess_inputs to DecoderInputPreparer.preprocess_inputs (#20134); the test still built a DecoderRunner, so every case failed with an AttributeError. Build the input preparer instead, with the overlap offsets under its buffers. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../executor/test_kv_lens_update_once.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/executor/test_kv_lens_update_once.py b/tests/unittest/_torch/executor/test_kv_lens_update_once.py index bee04e9db6df..f404a80713e1 100644 --- a/tests/unittest/_torch/executor/test_kv_lens_update_once.py +++ b/tests/unittest/_torch/executor/test_kv_lens_update_once.py @@ -1,19 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""RuntimeHooks.kv_lens_update_once: DecoderRunner._preprocess_inputs re-derives the KV-length metadata once, after the -overlap scheduler's KV-length correction, on exactly the steps that correct it, and leaves the inputs as before.""" +"""RuntimeHooks.kv_lens_update_once: DecoderInputPreparer.preprocess_inputs re-derives the KV-length metadata once, +after the overlap scheduler's KV-length correction, on exactly the steps that correct it, and leaves the inputs as +before.""" import types import pytest import torch -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks class _Metadata: - """The fields _preprocess_inputs reads; on_update_kv_lens records the kv_lens it saw.""" + """The fields preprocess_inputs reads; on_update_kv_lens records the kv_lens it saw.""" def __init__(self, gen: int, tokens_per_gen: int, with_kv_cache: bool = True): self.kv_cache_manager = object() if with_kv_cache else None @@ -29,12 +30,14 @@ def on_update_kv_lens(self) -> None: self.seen.append(self.kv_lens_cuda.tolist()) -def _runner(once: bool, disable_overlap: bool = False) -> DecoderRunner: - r = object.__new__(DecoderRunner) +def _preparer(once: bool, disable_overlap: bool = False) -> DecoderInputPreparer: + r = object.__new__(DecoderInputPreparer) r.runtime_hooks = RuntimeHooks(kv_lens_update_once=once) r._config = types.SimpleNamespace(disable_overlap_scheduler=disable_overlap) - r.previous_pos_id_offsets_cuda = torch.full((64,), -2, dtype=torch.int32) - r.previous_kv_lens_offsets_cuda = torch.full((16,), -3, dtype=torch.int32) + r._buffers = types.SimpleNamespace( + previous_pos_id_offsets_cuda=torch.full((64,), -2, dtype=torch.int32), + previous_kv_lens_offsets_cuda=torch.full((16,), -3, dtype=torch.int32), + ) r.mapping = types.SimpleNamespace(has_cp_helix=lambda: False) r.guided_decoder = None return r @@ -54,7 +57,7 @@ def test_corrected_step_updates_once_with_the_same_inputs(gen) -> None: outs = {} for once in (False, True): md = _Metadata(gen, t) - inputs = _runner(once)._preprocess_inputs( + inputs = _preparer(once).preprocess_inputs( _inputs(md, gen * t), enable_spec_decode=True, runtime_draft_len=t - 1 ) outs[once] = (md.seen, inputs["position_ids"].tolist(), md.kv_lens_cuda.tolist()) @@ -77,7 +80,7 @@ def test_uncorrected_steps_keep_their_update(spec, disable_overlap, with_kv_cach outs = {} for once in (False, True): md = _Metadata(2, 4, with_kv_cache) - inputs = _runner(once, disable_overlap)._preprocess_inputs( + inputs = _preparer(once, disable_overlap).preprocess_inputs( _inputs(md, 8), enable_spec_decode=spec, runtime_draft_len=3 ) outs[once] = (md.seen, inputs["position_ids"].tolist(), md.kv_lens_cuda.tolist()) 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 13/29] [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 ffa76abd2cdd04a5c798807870b017211a024f0b Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:58:12 -0700 Subject: [PATCH 14/29] [None][chore] MiniMax-M3 sm107 targets: kernels under ops, routing decline tests Review follow-ups mirrored from the sm103 targets, no kernel or numerics change: - Each sm107 target's kernels/ moves into its ops/ directory; imports and test references follow. - The targets' runtime-check docstrings no longer name the removed routing helper. - The routing test checks that the sm107 targets decline the same RoPE-base and sparse-layer layouts as the sm103 ones. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../kernels/__init__.py | 3 - .../minimax_m3_nvfp4__sm_107__tp2/modeling.py | 98 +++++++++---------- .../ops/attention/m3_fullctx_decode.py | 4 +- .../ops/attention/m3_index_score.py | 4 +- .../ops/attention/m3_select_long_dsl.py | 6 +- .../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_draft_lowm.py | 10 +- .../ops/gemm/m3_mx_gemv.py | 2 +- .../ops/gemm/m3_mx_gemv_bf16.py | 2 +- .../ops/gemm/m3_ugpu_head.py | 2 +- .../ops}/long_select.py | 2 +- .../{kernels => ops}/m3_attn_lc.py | 0 .../{kernels => ops}/m3_draft_step.py | 0 .../{kernels => ops}/m3_drafter_attn.py | 2 +- .../{kernels => ops}/m3_fold_tail.py | 0 .../{kernels => ops}/m3_force_vec.py | 0 .../{kernels => ops}/m3_fullctx_decode.py | 0 .../{kernels => ops}/m3_index_score.py | 0 .../{kernels => ops}/m3_locality.py | 0 .../{kernels => ops}/m3_lowm_direct.py | 0 .../{kernels => ops}/m3_lowm_splitk.py | 4 +- .../{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_moe_finalize.py | 0 .../{kernels => ops}/m3_prequant.py | 0 .../{kernels => ops}/m3_report.py | 0 .../{kernels => ops}/m3_router_gemv.py | 0 .../{kernels => ops}/m3_sampled_accept.py | 0 .../{kernels => ops}/m3_select_long_dsl.py | 0 .../{kernels => ops}/m3_shard_accept.py | 0 .../{kernels => ops}/m3_spec_tail.py | 0 .../{kernels => ops}/m3_target_head.py | 0 .../{kernels => ops}/m3_tokmap_guard.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 .../kernels/__init__.py | 3 - .../minimax_m3_nvfp4__sm_107__tp4/modeling.py | 72 +++++++------- .../ops/attention/m3_fullctx_decode.py | 4 +- .../ops/attention/m3_index_score.py | 4 +- .../ops/attention/m3_index_score_t2sc.py | 4 +- .../ops/attention/m3_select_long_dsl.py | 6 +- .../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_draft_lowm.py | 10 +- .../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_draft_lowm.py | 2 +- .../{kernels => ops}/m3_draft_step.py | 0 .../{kernels => ops}/m3_drafter_attn.py | 2 +- .../{kernels => ops}/m3_fold_tail.py | 0 .../{kernels => ops}/m3_force_vec.py | 0 .../{kernels => ops}/m3_fullctx_decode.py | 0 .../{kernels => ops}/m3_index_score.py | 0 .../{kernels => ops}/m3_index_score_t2sc.py | 2 +- .../{kernels => ops}/m3_lowm_splitk.py | 6 +- .../{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_report.py | 0 .../{kernels => ops}/m3_sampled_accept.py | 0 .../{kernels => ops}/m3_select_long_dsl.py | 0 .../{kernels => ops}/m3_shard_accept.py | 0 .../{kernels => ops}/m3_spec_tail.py | 0 .../{kernels => ops}/m3_target_head.py | 0 .../{kernels => ops}/m3_tokmap_guard.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 .../ops/norm/m3_embedding_firstnorm.py | 2 +- .../ops/spec/m3_force_vec.py | 2 +- .../ops/spec/m3_force_vec_req.py | 2 +- .../{kernels => ops}/triton_sparse_decode.py | 0 .../triton_sparse_decode_mq.py | 0 .../kernels/__init__.py | 3 - .../minimax_m3_nvfp4__sm_107__tp8/modeling.py | 90 ++++++++--------- .../ops/attention/m3_fullctx_decode.py | 4 +- .../ops/attention/m3_index_score.py | 4 +- .../ops/attention/m3_select_long_dsl.py | 6 +- .../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_draft_lowm.py | 8 +- .../ops/gemm/m3_mx_gemv.py | 2 +- .../ops/gemm/m3_mx_gemv_bf16.py | 2 +- .../ops}/long_select.py | 2 +- .../{kernels => ops}/m3_attn_lc8.py | 0 .../{kernels => ops}/m3_draft_step.py | 0 .../{kernels => ops}/m3_drafter_attn.py | 2 +- .../{kernels => ops}/m3_fc_wait.py | 0 .../{kernels => ops}/m3_fold_tail.py | 0 .../{kernels => ops}/m3_force_vec.py | 0 .../{kernels => ops}/m3_fullctx_decode.py | 0 .../{kernels => ops}/m3_index_score.py | 0 .../{kernels => ops}/m3_layer/__init__.py | 0 .../m3_layer/phase_producer.py | 0 .../{kernels => ops}/m3_locality.py | 0 .../{kernels => ops}/m3_lowm_splitk.py | 4 +- .../{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_report.py | 0 .../{kernels => ops}/m3_router_gemv.py | 0 .../{kernels => ops}/m3_sampled_accept.py | 0 .../{kernels => ops}/m3_select_long_dsl.py | 0 .../{kernels => ops}/m3_shard_accept.py | 0 .../{kernels => ops}/m3_spec_tail.py | 0 .../{kernels => ops}/m3_target_head.py | 0 .../{kernels => ops}/m3_tokmap_guard.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 ...est_modeling_v2_m3_block_table_contract.py | 4 +- .../test_modeling_v2_m3_drafter_fullctx.py | 6 +- ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 8 +- .../test_modeling_v2_m3_fullctx_decode.py | 6 +- .../test_modeling_v2_m3_n_valid_contract.py | 4 +- .../test_modeling_v2_m3_select_long_dsl.py | 4 +- .../test_modeling_v2_m3_sparse_decode.py | 4 +- .../test_modeling_v2_m3_sparse_decode_lc.py | 2 +- .../test_modeling_v2_m3_tokmap_guard.py | 4 +- .../test_modeling_v2_msa_kv_update.py | 6 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 4 +- .../comm/_m3_moe_push_op_matrix.py | 4 +- .../gemm/test_modeling_v2_m3_dense_seam.py | 6 +- .../gemm/test_modeling_v2_m3_ugpu_head.py | 2 +- .../gemm/test_r1b1_gemv_order_gate_a.py | 8 +- .../moe/test_modeling_v2_m3_moe_engine.py | 2 +- .../test_modeling_v2_m3_moe_finalize_rows.py | 4 +- .../moe/test_modeling_v2_m3_moe_flat_gate.py | 2 +- .../moe/test_modeling_v2_m3_moe_sh_take.py | 2 +- ...test_modeling_v2_m3_moe_triton_finalize.py | 4 +- .../test_modeling_v2_m3_capture_workspaces.py | 4 +- ...test_modeling_v2_m3_draft_ar_push_words.py | 6 +- .../spec/test_modeling_v2_m3_draft_input.py | 2 +- .../spec/test_modeling_v2_m3_draft_step.py | 10 +- .../spec/test_modeling_v2_m3_fold_tail.py | 4 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 10 +- .../test_modeling_v2_m3_sampled_accept.py | 6 +- .../test_modeling_v2_m3_dispatch.py | 4 +- .../test_modeling_v2_m3_rank_agreement.py | 2 +- .../test_modeling_v2_m3_wide_lc_cap.py | 2 +- .../test_modeling_v2_m3_drafter_fullctx.py | 8 +- ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 4 +- .../test_modeling_v2_m3_fullctx_decode.py | 8 +- .../test_modeling_v2_m3_n_valid_contract.py | 4 +- .../test_modeling_v2_m3_select_long_dsl.py | 4 +- .../test_modeling_v2_m3_sparse_decode.py | 8 +- .../test_modeling_v2_m3_tokmap_guard.py | 4 +- .../test_modeling_v2_msa_kv_update.py | 6 +- .../comm/_m3_capture_in_boundary_op_matrix.py | 4 +- .../comm/_m3_draft_ar_norm_op_matrix.py | 2 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 2 +- .../gemm/test_modeling_v2_m3_draft_lowm.py | 4 +- .../gemm/test_modeling_v2_m3_mx_gemv.py | 2 +- .../gemm/test_modeling_v2_m3_target_head.py | 2 +- .../test_modeling_v2_m3_moe_flat_gate_tp4.py | 2 +- .../spec/test_modeling_v2_m3_draft_input.py | 2 +- .../test_modeling_v2_m3_draft_push_take.py | 4 +- .../spec/test_modeling_v2_m3_draft_step.py | 2 +- ...t_modeling_v2_m3_draft_step_fixed_order.py | 6 +- .../test_modeling_v2_m3_exact_numerics.py | 4 +- .../spec/test_modeling_v2_m3_fold_tail.py | 4 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 12 +-- .../test_modeling_v2_m3_sampled_accept.py | 2 +- .../test_modeling_v2_m3_dispatch.py | 4 +- .../test_modeling_v2_m3_prebuild.py | 10 +- .../test_modeling_v2_m3_rank_agreement.py | 2 +- .../test_modeling_v2_m3_drafter_fullctx.py | 6 +- ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 4 +- .../test_modeling_v2_m3_fullctx_decode.py | 6 +- .../test_modeling_v2_m3_n_valid_contract.py | 4 +- .../test_modeling_v2_m3_select_long_dsl.py | 4 +- .../test_modeling_v2_m3_sparse_decode.py | 4 +- .../test_modeling_v2_m3_sparse_decode_lc8.py | 4 +- .../test_modeling_v2_m3_tokmap_guard.py | 4 +- .../test_modeling_v2_msa_kv_update.py | 6 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 4 +- .../comm/_m3_moe_push_op_matrix.py | 4 +- .../test_modeling_v2_m3_mx_gemv_producer.py | 2 +- .../test_modeling_v2_m3_swiglu_oai_bf16x2.py | 6 +- .../gemm/test_r1b1_gemv_order_gate_a.py | 8 +- .../test_modeling_v2_m3_capture_target.py | 2 +- .../test_modeling_v2_m3_capture_workspaces.py | 4 +- .../spec/test_modeling_v2_m3_di_fc_flag.py | 4 +- .../spec/test_modeling_v2_m3_draft_input.py | 2 +- .../spec/test_modeling_v2_m3_draft_step.py | 8 +- ...t_modeling_v2_m3_draft_step_fixed_order.py | 6 +- .../spec/test_modeling_v2_m3_fold_tail.py | 4 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 10 +- .../test_modeling_v2_m3_sampled_accept.py | 2 +- .../test_modeling_v2_m3_sampled_accept_1r.py | 2 +- .../test_modeling_v2_m3_dispatch.py | 6 +- .../test_modeling_v2_m3_rank_agreement.py | 2 +- .../modeling_v2/test_modeling_v2_routing.py | 33 +++++++ 229 files changed, 418 insertions(+), 406 deletions(-) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/embedding_firstnorm/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/embedding_firstnorm/m3_embedding_firstnorm_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/embedding_firstnorm/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8/kernels => minimax_m3_nvfp4__sm_107__tp2/ops}/long_select.py (98%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_attn_lc.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_draft_step.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_drafter_attn.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_fold_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_force_vec.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_fullctx_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_index_score.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_locality.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_lowm_direct.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_lowm_splitk.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_mc_allreduce.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_moe/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_moe/m3_moe_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_moe/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_moe_finalize.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_prequant.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_report.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_router_gemv.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_sampled_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_select_long_dsl.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_shard_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_spec_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_target_head.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/m3_tokmap_guard.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/msa_kv_update.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/mx_gemv/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/mx_gemv/m3_mx_gemv_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/mx_gemv/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/triton_sparse_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/{kernels => ops}/triton_sparse_decode_mq.py (100%) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/embedding_firstnorm/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp4/{kernels => ops}/embedding_firstnorm/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/long_select.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_attn_lc.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_draft_lowm.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_draft_step.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_drafter_attn.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_fold_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_force_vec.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_fullctx_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_index_score.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_index_score_t2sc.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_lowm_splitk.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_mc_allreduce.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_moe/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_moe/m3_moe16_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_moe/m3_moe_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_moe/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_moe/op16.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_prequant.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_report.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_sampled_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_select_long_dsl.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_shard_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_spec_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_target_head.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/m3_tokmap_guard.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/msa_kv_update.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/mx_gemv/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp4/{kernels => ops}/mx_gemv/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/{kernels => ops}/triton_sparse_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp8/kernels/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/embedding_firstnorm/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp8/{kernels => ops}/embedding_firstnorm/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2/kernels => minimax_m3_nvfp4__sm_107__tp8/ops}/long_select.py (98%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_attn_lc8.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_draft_step.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_drafter_attn.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_fc_wait.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_fold_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_force_vec.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_fullctx_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_index_score.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_layer/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_layer/phase_producer.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_locality.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_lowm_splitk.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_mc_allreduce.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_moe/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_moe/m3_moe_kernel.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_moe/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_prequant.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_report.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_router_gemv.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_sampled_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_select_long_dsl.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_shard_accept.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_spec_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_target_head.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/m3_tokmap_guard.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/msa_kv_update.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/mx_gemv/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp8/{kernels => ops}/mx_gemv/op.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/triton_sparse_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/{kernels => ops}/triton_sparse_decode_mq.py (100%) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/__init__.py deleted file mode 100644 index 3a3bd2995b75..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/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_107 / tp2 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py index 46329cb9d50d..2d02fac77dcf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py @@ -39,12 +39,12 @@ range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise ``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). A row range is empty unless a landed cell has certified it at the TP2 geometry on Rubin. Every fast path a GEN rank takes prints a rank-tagged -engagement line once per (path, rows) (``taken``). Kernel modules report through ``kernels/m3_report``, which this +engagement line once per (path, rows) (``taken``). Kernel modules report through ``ops/m3_report``, which this module wires to the same policy at import. **Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the -long-context decode for the o_proj. It publishes it in ``kernels/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +long-context decode for the o_proj. It publishes it in ``ops/m3_prequant.SLOT`` keyed by the BF16 tensor; the next MXFP8 projection on a decode graph takes it, and every take clears the slot. The decode boundaries' push targets (the o_proj / dense-down / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's exchange) are passed explicitly from the layer to the producer and back. @@ -166,8 +166,8 @@ 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 ( +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import ( m3_drafter_attn, m3_fold_tail, m3_mc_allreduce, @@ -175,15 +175,12 @@ m3_report, m3_spec_tail, ) -from .kernels import m3_fullctx_decode as m3_fullctx_decode_kernel -from .kernels import m3_index_score as _index_score_kernel -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 import triton_sparse_decode as _sparse_decode_kernel -from .kernels.long_select import long_select_blocks -from .kernels.m3_moe_finalize import m3_moe_finalize -from .kernels.mx_gemv import op as mx_gemv +from .ops import m3_fullctx_decode as m3_fullctx_decode_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_shard_accept as m3_shard_accept_kernel +from .ops import msa_kv_update as _msa_kv_update_kernel +from .ops import triton_sparse_decode as _sparse_decode_kernel from .ops.activation.m3_silu_and_mul import m3_silu_and_mul from .ops.attention.m3_draft_attention import m3_draft_attention from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope @@ -226,8 +223,11 @@ from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear from .ops.gemm.m3_ugpu_head import m3_ugpu_head +from .ops.long_select import long_select_blocks +from .ops.m3_moe_finalize import m3_moe_finalize from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm from .ops.spec.m3_capture_norm import m3_capture_norm from .ops.spec.m3_chunk_norm import m3_chunk_norm @@ -280,7 +280,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). # The multicast all-reduce + residual + RMSNorm at both boundaries of every TP2 decode step (one 8-token bucket), the # MoE boundary's also publishing the next QKV activation in MXFP8. SM107 TP2 hop: 3.5..4.4 us at 1..8 rows against # 8.7..9.5 us for the stock AUTO (NCCL) all-reduce + fused residual RMSNorm, and 13.3..14.0 us at the MoE boundary @@ -356,7 +356,7 @@ class DecodePlan: # TP2 skeleton: no cap (the built-in split) until a TP2 cap is measured and certified. DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) -# The dense layers' decode attention on PrimTS split-KV decode (kernels/m3_fullctx_decode) instead of the trtllm-gen +# The dense layers' decode attention on PrimTS split-KV decode (ops/m3_fullctx_decode) instead of the trtllm-gen # call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the # standalone reducer. Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 # (29.3), (1, 1) 15.7 (28.0), (4, 2) 24.5 (29.3), (1, 2) 20.2 (28.1). The stock split is capped at 16 cluster CTAs @@ -398,7 +398,7 @@ class DecodePlan: # The K/V length bound the full-context plans are compiled for; a step whose cache length bound is longer is a miss. It # covers the TP2 AgentX config's max_seq_len (1 << 20) plus the speculative rows the cache reserves past it. FULLCTX_DECODE_MAX_KV_LEN = 1 << 21 # capacity: fixed: the full-context plans' compiled K/V bound -# The Eagle3 drafter layer's generation phase on the same full-context decode (kernels/m3_drafter_attn: the in-tree +# The Eagle3 drafter layer's generation phase on the same full-context decode (ops/m3_drafter_attn: the in-tree # RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its # split at 8 CTAs per (request, KV head) on SM 10.7: 45-46 us per draft pass at 165k. DRAFTER_FULLCTX_ENABLED = True @@ -415,7 +415,7 @@ class DecodePlan: # thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request # multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by # that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, -# token-0 drafts). kernels/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# token-0 drafts). ops/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such # calls: what the build writes (padded rows included), so bitwise the op with a fresh map. False:. PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" @@ -490,7 +490,7 @@ class DecodePlan: # decode layer's MC_FINALIZE does. 48 as the agg canon runs it (v1.8.3; the same gates as above). GENERIC_MC_FINALIZE_MAX_TOKENS = 48 if MC_AR_TP2 else 0 # capacity: fixed: the finalize rows # The MoE finalize (+ the shared expert add) of a finalized MoE forward with at least MOE_TRITON_FINALIZE_MIN_TOKENS -# rows (prefill, mixed and wide decode past the folds) runs as kernels/m3_moe_finalize: the trtllm-gen runner returns +# rows (prefill, mixed and wide decode past the folds) runs as ops/m3_moe_finalize: the trtllm-gen runner returns # its FC2 output unfinalized and one Triton kernel gathers the top-k rows together, bitwise to the runner's own # finalize and the separate shared add (153 -> 98 us per layer at 16K rows on SM107). False: the runner finalizes. MOE_TRITON_FINALIZE = True @@ -513,7 +513,7 @@ class DecodePlan: # False: releases after the push. MOE_EARLY_AR_RELEASE = True # A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel -# over the target lm_head's BF16 shard logits (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 # B4 spec tail: the fused verify (shard / sampled accept), the fused first drafter inputs (chunk / capture norm) and the # fused draft loop (draft_step, its draft input, the drafter boundary all-reduce) also take the 9..32-row steps, up to @@ -524,7 +524,7 @@ class DecodePlan: B4_SPEC_TAIL = True SPEC_TAIL_MAX_TOKENS = DECODE_MAX_TOKENS if B4_SPEC_TAIL else TINY_MAX_TOKENS SPEC_TAIL_MAX_REQUESTS = SPEC_TAIL_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) -# W2 (SpecTail, agg wide decode): the greedy fused verify (kernels/m3_shard_accept) also takes the +# W2 (SpecTail, agg wide decode): the greedy fused verify (ops/m3_shard_accept) also takes the # 33..WIDE_MAX_TOKENS-row decode-only steps, as whole-request chunks of the kernel's certified 48-row geometry (12 # requests per launch at draft length 3; m3_shard_accept.WIDE_SPEC_TAIL), instead of the gathered target logits + # argmax; and the in-tree draft loop takes the drafter's top-1 from its local head shard (W2b). Sampled verifies keep @@ -538,7 +538,7 @@ class DecodePlan: # with the coded head 117.0 us; the BF16-head + shard-accept arm is not measured yet. SHARD_ACCEPT_MAX_REQUESTS = WIDE_VERIFY_MAX_REQUESTS # A sampled (temperature + top-p, advanced_sampling_mode no_topk) Eagle3 verify step of up to this many requests -# verifies with the fused sampled acceptance (kernels/m3_sampled_accept, the TP2-shard kernel) over the target +# verifies with the fused sampled acceptance (ops/m3_sampled_accept, the TP2-shard kernel) over the target # lm_head's BF16 shard logits, at the configured draft length (the AgentX contract). 0: (the gathered logits # and the in-tree sampler run). Rubin TP2 b1/b2 (p5 microbench): 180.1 / 181.8 -> 115.1 / 113.9 us per verify. SAMPLED_VERIFY_MAX_REQUESTS = SPEC_TAIL_MAX_REQUESTS @@ -558,7 +558,7 @@ class DecodePlan: # to this many verify rows; wider first passes take the in-tree first drafter inputs. 0:. CHUNK_NORM_MAX_ROWS = SPEC_TAIL_MAX_TOKENS # The fused draft loop's two all-reduce boundaries per pass (o_proj -> post-attention norm, down -> final norm) of up -# to this many rows run the multicast all-reduce + residual + RMSNorm (kernels/m3_mc_allreduce, opts bit 16: the +# to this many rows run the multicast all-reduce + residual + RMSNorm (ops/m3_mc_allreduce, opts bit 16: the # drafter's stock rounding, normalizing the unrounded FP32 sum as FlashInfer fused_add_rmsnorm does) instead of the # built-in all-reduce + RMSNorm. 0:. Rubin TP2 (p16 microbench): 10.0..11.0 -> 4.6..5.3 us per boundary. # (The drafter epilogue, opts bit 16, is built for the 8-token bucket only: a B4 spec-tail first pass of 16 / 32 rows @@ -573,7 +573,7 @@ class DecodePlan: # The drafter boundary all-reduce fed by its projection's push (DRAFT_AR_PUSH): its own taken / not-taken lines. DRAFT_AR_PUSH_PATH = "m3_draft_ar_push" # The same push from the drafter's direct ('d') low-M linears (the 1-row passes' o_proj and down at TP2), through this -# target's copy of FlashInfer's direct kernel (kernels/m3_lowm_direct; its output FlashInfer's): the owning thread +# target's copy of FlashInfer's direct kernel (ops/m3_lowm_direct; its output FlashInfer's): the owning thread # stores the CTA's column pairs into the exchange before the kernel's last-instruction release. Bitwise. Rubin TP2 # world 2 (dpush-1 782685), per 1-row boundary: o_proj 9.22 -> 8.09 us, down 13.81 -> 12.72. False:. DRAFT_AR_PUSH_DIRECT = True @@ -593,9 +593,9 @@ class DecodePlan: # late. 0:. Rubin TP2 c1 (locked nsys 768559): the head started 17.4 us after its input was ready. FC_AFTER_HEAD_MAX_ROWS = SPEC_TAIL_MAX_TOKENS # A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (spec_worker): per pass the -# draft input in one launch (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. @@ -608,7 +608,7 @@ class DecodePlan: # keep the in-tree cuBLAS GEMM. 0:. Rubin TP2 (mclk locked, 4 rotating sets) at 1 row: qkv 13.18 -> 9.71, # o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the certified low-M rows -# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (kernels/m3_lowm_splitk): bit for bit the +# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (ops/m3_lowm_splitk): bit for bit the # same output at every certified cell (spec/test_modeling_v2_m3_lowm_splitk.py); the seam for the drafter-boundary push # epilogue and the down projection's SwiGLU prologue. False: (FlashInfer's kernel). DRAFT_LOWM_OWN_SPLITK = True @@ -617,7 +617,7 @@ class DecodePlan: for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() } # The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows -# (kernels/triton_sparse_decode.MQ_MIN_TOKENS); the multi-query producer needs one KV head, so at the TP2 heads every +# (ops/triton_sparse_decode.MQ_MIN_TOKENS); the multi-query producer needs one KV head, so at the TP2 heads every # identity step takes the per-row producer. SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover SPARSE_DECODE_KV_HEADS = 2 # capacity: fixed: the tp2 KV heads per rank @@ -640,7 +640,7 @@ class DecodePlan: # 16-CTA cluster launch and prologue then overlap the rank: Rubin TP2 in-model, locked memory clock, AgentX draws, # 2 instances vs the r4.3 stored A: replay period -162 us fixed, -94 at a 323K-token request. SELECT_LONG_TRIGGER_PHASE = 1 -# The long-row select runs as kernels/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per +# The long-row select runs as ops/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per # (query row, KV head) instead of one 1024-thread CTA, the same selection bit for bit (ties included), with 4, 11 or # 24 keys per thread (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices. Rubin TP2 decode chain (2 KV heads, # 4 rows, 57 layers behind the producer stream and the scorer, both orders x 3 rounds): -0.127 / -0.148 / -0.037 us @@ -683,10 +683,10 @@ class DecodePlan: # sequence). FORCE_AL_PER_REQUEST = True -# The decode MXFP8 GEMVs (kernels/mx_gemv, SM107 TP2 table), per projection (N, K): the rows each serves, on an +# The decode MXFP8 GEMVs (ops/mx_gemv, SM107 TP2 table), per projection (N, K): the rows each serves, on an # activation its producer quantized (the MoE boundary for the QKV, the attention merge for the o_proj) and on a bf16 # one (quantized in the GEMV's prologue, bit-identical to mxfp8_quantize). SM107 graph us per projection at 1..8 rows: -# 2.2..8.6 against 4.7..10.5 for the stock MXFP8 quantize + FlashInfer CuTe DSL GEMM (kernels/mx_gemv/op.py). +# 2.2..8.6 against 4.7..10.5 for the stock MXFP8 quantize + FlashInfer CuTe DSL GEMM (ops/mx_gemv/op.py). # False puts every projection back on the stock decode-graph GEMM. # On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the # built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. @@ -1284,7 +1284,7 @@ def reset_logged() -> None: _logged.clear() -# Kernel modules report from inside their calls through kernels/m3_report: the same policy. +# Kernel modules report from inside their calls through ops/m3_report: the same policy. m3_report.install( on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed ) @@ -1383,8 +1383,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. @@ -1805,7 +1805,7 @@ def _decode_routed_unfinalized(core, i: int, h: torch.Tensor): def _triton_finalize(rows: int, routed, shared: Optional[torch.Tensor]) -> torch.Tensor: - """The finalized routed output (+ ``shared``) from the runner's unfinalized parts via kernels/m3_moe_finalize + """The finalized routed output (+ ``shared``) from the runner's unfinalized parts via ops/m3_moe_finalize (``MOE_TRITON_FINALIZE``). Raises if the parts are not the trtllm-gen layout (the runner did not finalize them: no fallback).""" out = None @@ -2679,9 +2679,9 @@ def prebuild_kernels(self) -> Dict[str, object]: ) if not self.mxfp8.quant_method.use_cutlass: raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") - from .kernels import m3_fullctx_decode as _fullctx_decode_kernel - from .kernels import m3_select_long_dsl as _select_long_dsl_kernel - from .kernels import m3_tokmap_guard + from .ops import m3_fullctx_decode as _fullctx_decode_kernel + from .ops import m3_select_long_dsl as _select_long_dsl_kernel + from .ops import m3_tokmap_guard device = self.w["embed"].device # The one-request token-map guard of the stock generation preprocess (``PREPROCESS_TOKMAP_FIX``): every path @@ -4048,7 +4048,7 @@ def _dense_attention(self, i: int, position_ids, hidden, step, push): def _fullctx_plan(self, i: int, device: torch.device, step): """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs - on ``kernels/m3_fullctx_decode`` (``FULLCTX_DECODE``): a covered (query rows per request, requests) shape over + on ``ops/m3_fullctx_decode`` (``FULLCTX_DECODE``): a covered (query rows per request, requests) shape over the layer's FP8 pool. None otherwise (the MSA backend runs). Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 (29.3), (1, 1) @@ -4073,7 +4073,7 @@ def _fullctx_plan(self, i: int, device: torch.device, step): if pool.dtype != torch.float8_e4m3fn: missed(path, num_tokens, "the dense layer's KV pool is not FP8") return None - from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + from .ops import m3_fullctx_decode as _fullctx_decode_kernel own_reducer = FULLCTX_DECODE_OWN_REDUCER if own_reducer and not _fullctx_decode_kernel.has_own_reducer(): @@ -4095,7 +4095,7 @@ def _fullctx_decode(self, i: int, q, k, v, step, plan_tokens, *, kv_written: boo write_msa_phase_kv, ) - from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + from .ops import m3_fullctx_decode as _fullctx_decode_kernel core = self.core plan, num_tokens = plan_tokens @@ -4264,7 +4264,7 @@ def _select(self, i: int, max_score, step, num_tokens: int): if group > 1 else max_score ) - from .kernels import m3_select_long_dsl as _select_long_dsl_kernel + from .ops import m3_select_long_dsl as _select_long_dsl_kernel num_blocks = int(max_score_kv.shape[1]) table = None @@ -4321,7 +4321,7 @@ def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: selected blocks recur one step later; same layer, not the next one as at TP4: the MoE between two sparse layers streams ~250 MB through the 126 MB L2, so an earlier issue would be evicted; hints only). """ - from .kernels import m3_attn_lc as _attn_lc_kernel + from .ops import m3_attn_lc as _attn_lc_kernel core = self.core backend = core.attn[i] @@ -4633,10 +4633,10 @@ class M3Eagle3Draft(nn.Module): (``ops/gemm/m3_draft_lowm``, the TP8 port) at the rows ``DRAFT_LOWM`` covers, compiled before the warmup (``prebuild_linears``); every other call is the in-tree cuBLAS GEMM. Rubin TP2 (mclk locked, 4 rotating sets) at 1 row: qkv 13.18 -> 9.71, o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. The - split-K ('s') tactics run this target's copy of FlashInfer's kernel (``kernels/m3_lowm_splitk``, bit for bit the + split-K ('s') tactics run this target's copy of FlashInfer's kernel (``ops/m3_lowm_splitk``, bit for bit the same output); with ``DRAFT_AR_PUSH`` the o_proj and down projection of the fused draft loop also push their rows into the drafter boundary all-reduce's exchange (``arm_push`` / ``take_pushed``), as do the direct ('d') 1-row - linears through this target's copy of the direct kernel (``kernels/m3_lowm_direct``; o_proj 9.22 -> 8.09 us, + linears through this target's copy of the direct kernel (``ops/m3_lowm_direct``; o_proj 9.22 -> 8.09 us, down 13.81 -> 12.72 us per 1-row boundary). """ @@ -4809,8 +4809,8 @@ def _lowm_linear(self, weight: torch.Tensor, x: torch.Tensor) -> Optional[torch. def prebuild_linears(self) -> int: """Compile (and run once) every covered row count of every drafter projection, outside capture, on the current stream, before the executor's warmup. Returns the number of calls compiled.""" - from .kernels.m3_lowm_direct import compile_direct_push - from .kernels.m3_lowm_splitk import compile_push + from .ops.m3_lowm_direct import compile_direct_push + from .ops.m3_lowm_splitk import compile_push use_pdl = get_env_enable_pdl() built = 0 @@ -4830,7 +4830,7 @@ def prebuild_linears(self) -> int: # (its pushes need the all-reduce's exchange and a matching call). compile_push(x, weight.t(), out, use_pdl, tactic) if DRAFT_AR_PUSH and DRAFT_AR_PUSH_DIRECT and kind == "d" and tactic[1] % 2 == 0: - # The direct kernel's push variant (kernels/m3_lowm_direct), compiled only. + # The direct kernel's push variant (ops/m3_lowm_direct), compiled only. compile_direct_push(x, weight.t(), out, use_pdl, tactic) self._lowm_built.add((n, k, m, use_pdl)) built += 1 @@ -4887,7 +4887,7 @@ def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is False, where the caller fuses the reduction). The backend's generation phase runs this target's full-context - decode with the fused RoPE + FP8 K/V append when the shell installed it (``kernels/m3_drafter_attn``).""" + decode with the fused RoPE + FP8 K/V append when the shell installed it (``ops/m3_drafter_attn``).""" qkv = self.linear(x, "qkv") _record("m3_draft_attention", int(md.num_tokens)) if is_in_breakable_cuda_graph(): @@ -6023,7 +6023,7 @@ def greedy_sample_draft_with_tp_gather( self, logits: torch.Tensor, spec_metadata=None, mapping_lm_head_tp=None ) -> torch.Tensor: """W2b (``WIDE_SPEC_TAIL``): the in-tree draft loop's greedy top-1 over this rank's vocab shard of the drafter - logits, from the fused shard argmax + cross-rank key exchange (kernels/m3_shard_accept at draft length 0, at + logits, from the fused shard argmax + cross-rank key exchange (ops/m3_shard_accept at draft length 0, at most 48 rows per launch) instead of the per-rank max, the all-gather of (index, value) pairs and the pick. The same tokens (keys order BF16 values exactly; ties take the lowest global index, as the pick over ranks and the first-index max within a shard). Plain TP only (no LM-head TP in ADP); else, or when the logits are not a diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py index f79da95e389c..bc3cd23b289f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The dense layers' full-context decode attention at TP2 on Rubin (``kernels/m3_fullctx_decode``): PrimTS split-KV +"""The dense layers' full-context decode attention at TP2 on Rubin (``ops/m3_fullctx_decode``): PrimTS split-KV decode at the recipe modeling's ``FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" from __future__ import annotations @@ -12,7 +12,7 @@ fp8_e4m3_saturate, ) -from ...kernels import m3_fullctx_decode as kernel +from .. import m3_fullctx_decode as kernel from .m3_dense_decode import dense_attention_reference, gather_paged_kv __all__ = ["m3_fullctx_decode"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py index e9ed36eec6f8..be1c89b97261 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py @@ -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 @@ -45,7 +45,7 @@ def _wide_lens(requests: int) -> tuple: 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_107__tp2/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py index 191b4e3132ae..205e304d5f8e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py @@ -12,7 +12,7 @@ class M3SelectLongDsl(OpWrapper): - """``m3_select_long``'s contract, cells and reference, run by ``kernels/m3_select_long_dsl``: the same selection + """``m3_select_long``'s contract, cells and reference, run by ``ops/m3_select_long_dsl``: the same selection bit for bit as the 1024-thread ``minimax_m3_sm107_tp2_select_long`` (also on rows where more than 2048 blocks tie at the 16th-largest score), from one 384-thread CTA per row. Returns ``out``.""" @@ -105,7 +105,7 @@ def is_valid( The kernel is compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; the target compiles its keys before warmup - (`kernels/m3_select_long_dsl.prebuild`), and a call outside them compiles + (`ops/m3_select_long_dsl.prebuild`), and a call outside them compiles on first use. """ @@ -120,7 +120,7 @@ def __call__( early_trigger: bool = False, trigger_phase: int = 0, ) -> torch.Tensor: - from ...kernels.m3_select_long_dsl import select_long_dsl + from ..m3_select_long_dsl import select_long_dsl if scores.shape[1] > MAX_BLOCKS: raise RuntimeError( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py index 0a8613019283..98f04538e784 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py @@ -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_107__tp2/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py index 883f38e0803d..b624ae64e59e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py @@ -97,7 +97,7 @@ class M3McAllreduceNorm(OpWrapper): ARCHS = frozenset({Arch.SM_107}) # The tp2 decode boundaries (hidden 6144, 2 ranks, eps 1e-6): every decode step fits the 8-token bucket, epilogue - # option 7 at 1, 2, 7, 8 tokens and 15 at 3..6 (kernels/m3_mc_allreduce.AR_OPTS_BY_M; every option is + # option 7 at 1, 2, 7, 8 tokens and 15 at 3..6 (ops/m3_mc_allreduce.AR_OPTS_BY_M; every option is # bit-identical). CELLS: tuple[Cell, ...] = ( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/embedding_firstnorm/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py index 427ee2598f91..94f5e112e592 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py @@ -99,7 +99,7 @@ def __call__( if kind == "d" and ar is not None: # The boundary push (modeling's DRAFT_AR_PUSH_DIRECT): this target's copy of the direct kernel, whose # output is FlashInfer's, also stores the rows into the boundary all-reduce's exchange. - from ...kernels.m3_lowm_direct import DirectTactic, run_direct_dense_push + from ..m3_lowm_direct import DirectTactic, run_direct_dense_push run_direct_dense_push(x, weight.t(), out, use_pdl, DirectTactic(*tactic), ar) elif kind == "d": @@ -110,13 +110,9 @@ def __call__( run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) elif own_splitk: - # This target's copy of the split-K kernel (kernels/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` # it also pushes the rows into a drafter boundary all-reduce (modeling's DRAFT_AR_PUSH). - from ...kernels.m3_lowm_splitk import ( - SplitKTactic, - run_splitk_dense, - run_splitk_dense_push, - ) + from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push if ar is not None: run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py index 6f04ed1ba6d4..96352e76d700 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py @@ -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", "served_max_rows"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py index 5788c113d449..0398fe64b3d5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py @@ -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_107__tp2/ops/gemm/m3_ugpu_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py index 335effba7fee..c734b2a3172e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py @@ -75,7 +75,7 @@ def prepare_head(lm_head: torch.nn.Module) -> Optional[UgpuHead]: Before capture.""" from tensorrt_llm._torch.modules import low_m_gemm - from ...kernels.m3_locality import LocalityMap + from ..m3_locality import LocalityMap w = getattr(lm_head, "weight", None) if ( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/long_select.py similarity index 98% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/long_select.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/long_select.py index 14e24fbdbcac..79aaf6674300 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/long_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/long_select.py @@ -14,7 +14,7 @@ The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long`` (this target's copy of ``minimax_m3_select_long`` with a choice of where it releases the sparse decode), called through this target's ``m3_select_long`` entry, or with -``dsl`` the same selection from one 384-thread CTA per row (``m3_select_long_dsl``, kernels/m3_select_long_dsl). +``dsl`` the same selection from one 384-thread CTA per row (``m3_select_long_dsl``, ops/m3_select_long_dsl). """ from typing import Optional diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_attn_lc.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_attn_lc.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_attn_lc.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_draft_step.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_drafter_attn.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py index eabe675e6970..a4c7b5ccb8fd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_drafter_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py @@ -8,7 +8,7 @@ This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context -PrimTS split-KV decode of ``kernels/m3_fullctx_decode`` at the same pinned recipes as the dense layers. +PrimTS split-KV decode of ``ops/m3_fullctx_decode`` at the same pinned recipes as the dense layers. With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead (``minimax_m3_sm107tp2_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fold_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_force_vec.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_force_vec.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_force_vec.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_fullctx_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_index_score.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_locality.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_direct.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_direct.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_direct.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_direct.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_splitk.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py index 466a327a4bbf..027541ea6990 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_lowm_splitk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py @@ -317,7 +317,7 @@ def _store_shared_remote_v4( # The drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH): the exchange words of -# kernels/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +# ops/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. @dsl_user_op def _ld_volatile_s32(addr, *, loc=None, ip=None): return Int32( @@ -935,7 +935,7 @@ def _push_tile( ar_n: cutlass.Int32, ): """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output - would store them (kernels/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + would store them (ops/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" cute_ext.partition_and_copy(thr_t2r, rD, sD) cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_mc_allreduce.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/m3_moe_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe_finalize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe_finalize.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_moe_finalize.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe_finalize.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_prequant.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_report.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_report.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_report.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_router_gemv.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_sampled_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_sampled_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_sampled_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_select_long_dsl.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_shard_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_spec_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_spec_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_spec_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_target_head.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/m3_tokmap_guard.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py index 562a9f392cef..6e3feee8882a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py @@ -17,7 +17,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_107__tp2/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/msa_kv_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/m3_mx_gemv_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/mx_gemv/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py index 9093d53b3ce7..94827619dad2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py @@ -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_107__tp2/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py index 2001bf183b54..a192e89c5c76 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py @@ -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, force_vec_req_kernel +from ..m3_force_vec import force_vec_kernel, force_vec_req_kernel class M3ForceVec(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/triton_sparse_decode_mq.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/__init__.py deleted file mode 100644 index 91d4ff843d5e..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107 / tp4 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py index 997280727b01..fbc60ddc1d76 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py @@ -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 -long-context decode for the o_proj. It publishes it in ``kernels/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +long-context decode for the o_proj. It publishes it in ``ops/m3_prequant.SLOT`` keyed by the BF16 tensor; the next MXFP8 projection on a decode graph takes it (``_linear``), and every take clears the slot. The decode boundaries' push targets (the o_proj / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's exchange) are passed explicitly from the layer to the producer and back. @@ -155,22 +155,19 @@ from tensorrt_llm.quantization.mode import QuantAlgo from . import weights as _weights -from .kernels import m3_attn_lc, m3_fold_tail, m3_mc_allreduce, m3_prequant, m3_report, m3_spec_tail -from .kernels import m3_draft_lowm as _draft_lowm_kernel -from .kernels import m3_draft_step as m3_draft_step_kernel -from .kernels import m3_drafter_attn as _drafter_attn_kernel -from .kernels import m3_fullctx_decode as _fullctx_decode_kernel -from .kernels import m3_index_score as _index_score_kernel -from .kernels import m3_index_score_t2sc as _index_score_t2sc_kernel -from .kernels import m3_sampled_accept as m3_sampled_accept_kernel -from .kernels import m3_select_long_dsl as _select_long_dsl_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 m3_tokmap_guard as _tokmap_guard -from .kernels import msa_kv_update as _msa_kv_update_kernel -from .kernels.long_select import long_select_blocks -from .kernels.mx_gemv import op as mx_gemv -from .kernels.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS +from .ops import m3_attn_lc, m3_fold_tail, m3_mc_allreduce, m3_prequant, m3_report, m3_spec_tail +from .ops import m3_draft_lowm as _draft_lowm_kernel +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import m3_drafter_attn as _drafter_attn_kernel +from .ops import m3_fullctx_decode as _fullctx_decode_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_index_score_t2sc as _index_score_t2sc_kernel +from .ops import m3_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_select_long_dsl as _select_long_dsl_kernel +from .ops import m3_shard_accept as m3_shard_accept_kernel +from .ops import m3_target_head as m3_target_head_kernel +from .ops import m3_tokmap_guard as _tokmap_guard +from .ops import msa_kv_update as _msa_kv_update_kernel from .ops.activation.m3_silu_and_mul import m3_silu_and_mul from .ops.attention.m3_dense_decode import m3_dense_decode from .ops.attention.m3_dense_qk_norm_rope_fp8 import m3_dense_qk_norm_rope_fp8 @@ -214,9 +211,11 @@ from .ops.gemm.m3_mxfp8_pinned_quantize import m3_mxfp8_pinned_quantize from .ops.gemm.m3_target_head import MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS from .ops.gemm.m3_target_head import m3_target_head +from .ops.long_select import long_select_blocks from .ops.moe.m3_moe16_engine import M3Moe16State, m3_moe16_engine from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm from .ops.spec.m3_draft_input import m3_draft_input from .ops.spec.m3_draft_silu import m3_draft_silu @@ -227,6 +226,7 @@ from .ops.spec.m3_shard_accept import m3_shard_accept from .ops.spec.m3_spec_offsets import m3_spec_offsets from .ops.spec.m3_step_prologue import m3_step_prologue +from .ops.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS _SM = (10, 7) # capacity: fixed: the device capability, not a size # -------------------------------------------------------------------------------------------------------------------- @@ -252,7 +252,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). # SM 10.7 TP4: one bucket. MC_AR_TP4 puts every decode boundary on the multicast all-reduce + residual + RMSNorm # (trtllm::minimax_m3_sm107_tp4_mc_allreduce_norm, its row kernel: bitwise to the stock TP4 oneshot all-reduce fusion # on the residual and the normed rows) with the MXFP8 copy for the next QKV; False = the built-in all-reduce. @@ -286,7 +286,7 @@ class DecodePlan: # 4.79 -> 4.14 at 4 rows, ~ -2..-3.5 us/step. Kill: False (the row kernel, bitwise the stock boundary). # OFF: in-model timing was flat and natural acceptance length was grey (lower bound -1.16% vs the -1% line). DRAFT_MC_AR_CLUSTER = False -# The drafter's o_proj and down projection, on the copied split-K kernel (kernels/m3_lowm_splitk), also push their rows +# The drafter's o_proj and down projection, on the copied split-K kernel (ops/m3_lowm_splitk), also push their rows # into the drafter boundary all-reduce's exchange (the row kernel's slots: the canonical words its own push would store, # at the (token, 768-column slice) call parity), which then runs prepushed 2 instead of pushing x itself; the copy # releases its dependents only after its own grid wait and the pushes (the prepushed >= 2 contract). Bitwise. The TP8 @@ -321,7 +321,7 @@ class DecodePlan: # TP4 skeleton: no cap (the built-in split) until a TP4 SM 10.7 cap is measured and certified. DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) -# The dense layers' decode attention on PrimTS split-KV decode (kernels/m3_fullctx_decode) instead of the trtllm-gen +# The dense layers' decode attention on PrimTS split-KV decode (ops/m3_fullctx_decode) instead of the trtllm-gen # call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the # BF16 merge. At TP4 (16 query heads, 1 KV head per rank) the stock split is capped at 16 cluster CTAs per (request, # KV head), i.e. 16 of 212 CTAs at one request: stock dense fmhaSm107a 33.4-33.8 us per call and the stock drafter @@ -380,7 +380,7 @@ class DecodePlan: FULLCTX_DECODE_MAX_KV_LEN = ( 1 << 21 ) # capacity: fixed: the plans' KV-length bound (1M context + spec rows) -# The Eagle3 drafter layer's generation phase on the same full-context decode (kernels/m3_drafter_attn: the in-tree +# The Eagle3 drafter layer's generation phase on the same full-context decode (ops/m3_drafter_attn: the in-tree # RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its # split at 8 CTAs per (request, KV head) on SM 10.7: 45.4 us per draft pass at 165k. Kill constant. DRAFTER_FULLCTX_ENABLED = True @@ -393,7 +393,7 @@ class DecodePlan: # thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request # multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by # that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, -# token-0 drafts). kernels/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# token-0 drafts). ops/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such # calls: what the build writes (padded rows included), so bitwise the op with a fresh map. Kill: False. PREPROCESS_TOKMAP_FIX = True PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" @@ -451,13 +451,13 @@ class DecodePlan: MC_FINALIZE_MIN_TOKENS = 9 # capacity: fixed: the rows above the engines MC_FINALIZE_MAX_TOKENS = 0 # capacity: fixed: off at TP4 # A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel -# over the target lm_head's BF16 shard logits (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 sampled (temperature + top-p, advanced_sampling_mode=no_topk) Eagle3 verify step of up to this many requests -# verifies with the fused sampled acceptance (kernels/m3_sampled_accept) over this rank's logits shard: every request +# verifies with the fused sampled acceptance (ops/m3_sampled_accept) over this rank's logits shard: every request # count a decode step holds. SAMPLED_VERIFY_MAX_REQUESTS = SHARD_ACCEPT_MAX_REQUESTS # capacity: fixed: a decode step's requests # The sampled verify runs the sm_107 tp4 copy of the TP2 sampled-accept kernel @@ -496,7 +496,7 @@ class DecodePlan: # gather) up to this many requests: every request count a decode step holds. FIRST_DRAFTER_INPUTS_MAX_REQUESTS = SHARD_ACCEPT_MAX_REQUESTS # capacity: fixed: step requests # A fused-verify step of up to this many verify rows (one or two requests at draft length 3) computes the target -# logits shard from the 12-bit coded copy of the BF16 head (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. # Off at TP4 (as at TP2): at the TP4 shard (50,016 x 6144) the in-tree BF16 cuBLAS head is faster than the coded kernel, # 47.9 vs 56.8 us at 4 and 8 rows (TP2: 93 vs 109). True re-enables the coded head. @@ -512,9 +512,9 @@ class DecodePlan: # holds the SMs first and the head GEMM starts late (TP2: 17.4 us late). Kill: 0. FC_AFTER_HEAD_MAX_ROWS = 8 # capacity: fixed: measured fork-after-head crossover # A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (M3Eagle3Worker): per pass the -# draft input in one launch (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. @@ -547,7 +547,7 @@ class DecodePlan: for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() } # The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows -# (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. @@ -599,7 +599,7 @@ class DecodePlan: # warp-0 tail runs ~0.95 us/layer slower in-model next to the early-launched sparse decode (+53 us/step); # on at TP4 only after an in-model A,B shows a gain. SELECT_LONG_SL3 = False -# The long-row select runs as kernels/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per row +# The long-row select runs as ops/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per row # instead of one 1024-thread CTA, the same selection bit for bit, with 4, 11 or 24 keys per thread # (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices (SELECT_LONG_SL3 does not apply). TP4 has the TP8 target's # per-rank select geometry (1 KV head, 4 rows per request), where on the Rubin decode chain (57 layers behind the @@ -617,7 +617,7 @@ class DecodePlan: # 1 / 4 / 8. Timing only: the tiles are read after the wait. Kill: 1. INDEX_SCORE_EARLY_MAX_REQUESTS = 2 # capacity: fixed: measured early-issue crossover # T2SC: up to 2 requests the sparse decode scores with this target's copy of the sm_107 TP2 / TP8 scorer -# (kernels/m3_index_score_t2sc; ops/attention/m3_index_score_t2sc), whose split is chosen per request count and which +# (ops/m3_index_score_t2sc; ops/attention/m3_index_score_t2sc), whose split is chosen per request count and which # TMA-loads an all-old first block before its grid wait; the scores are bit-identical # to the in-tree scorer. TP4 has the TP8 target's per-rank indexer geometry (1 index head over the replicated index-K # of its KV group), where this schedule is ~100 us per decode step faster at 1 request than split 512 without the @@ -635,7 +635,7 @@ class DecodePlan: # (the MoE boundary for the QKV, the attention merge for the o_proj) the GEMV serves every row up to 48 but the # unsplit K=768 shared down projection at 17..32 rows, where FlashInfer's GEMM is faster. # The QKV MXFP8 GEMVs' (split, ring) at decode rows up to GEMV_QKV_CFG_MAX_ROWS: a TP4-owned override of the shared -# GB300-tuned table (kernels/mx_gemv/op.py). SM107 TP4 sweep +# GB300-tuned table (ops/mx_gemv/op.py). SM107 TP4 sweep # (4 rows): sparse-layer QKV (8, 8) 2.70 us vs (4, 12) 3.01, dense-attention QKV (8, 8) 2.60 vs 2.83. Split 8 # sums K in a different fp32 split-K order than split 4: exact-class, not bitwise to the shared config. Kill: False. GEMV_QKV_SPLIT8_TP4 = True @@ -1022,7 +1022,7 @@ def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: # at its tail (one launch inside the step graph), so neither the sampler's store nor the next step's gathers run # between the step graphs (M3SpecStepFold.in_forward). STEP_SPEC_FOLD_IN_FORWARD = True -# The forward's tail (STEP_SPEC_FOLD_IN_FORWARD) runs as this target's C++ op (kernels/m3_fold_tail cpp=, +# The forward's tail (STEP_SPEC_FOLD_IN_FORWARD) runs as this target's C++ op (ops/m3_fold_tail cpp=, # minimaxM3Sm107Tp4/foldTail.cu) instead of the Triton kernel: the same integer copies, one launch either way (no # Triton on the generation decode path). Kill: False. STEP_FOLD_TAIL_CPP = True @@ -1227,7 +1227,7 @@ def reset_logged() -> None: _logged.clear() -# Kernel modules report from inside their calls through kernels/m3_report: the same policy. +# Kernel modules report from inside their calls through ops/m3_report: the same policy. m3_report.install( on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed ) @@ -1326,8 +1326,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. @@ -3702,7 +3702,7 @@ def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = def _fullctx_plan(self, i: int, device: torch.device, step): """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs on - kernels/m3_fullctx_decode: a covered (query rows per request, requests) shape over the layer's FP8 pool. None + ops/m3_fullctx_decode: a covered (query rows per request, requests) shape over the layer's FP8 pool. None otherwise (the MSA backend runs).""" path = FULLCTX_DECODE.path if is_torch_compiling(): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py index ad3abcadaae9..2c6c2133b2bb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The dense layers' full-context decode attention at TP4 on Rubin (``kernels/m3_fullctx_decode``): PrimTS split-KV +"""The dense layers' full-context decode attention at TP4 on Rubin (``ops/m3_fullctx_decode``): PrimTS split-KV decode at the recipe ``modeling.FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" from __future__ import annotations @@ -12,7 +12,7 @@ fp8_e4m3_saturate, ) -from ...kernels import m3_fullctx_decode as kernel +from .. import m3_fullctx_decode as kernel from .m3_dense_decode import dense_attention_reference, gather_paged_kv __all__ = ["m3_fullctx_decode"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py index f44d566310e7..4fb706182485 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py @@ -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_107__tp4/ops/attention/m3_index_score_t2sc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py index 45103a9e289c..39a27ee64b71 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py @@ -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_t2sc`` (T2SC: the sm_107 +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score_t2sc`` (T2SC: the sm_107 TP2 / TP8 scorer as this target's copy, at its 1 index head per rank; ``modeling.INDEX_SCORE_T2SC``).""" from __future__ import annotations @@ -27,7 +27,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_t2sc``).""" - from ...kernels.m3_index_score_t2sc import IndexScoreRunner + from ..m3_index_score_t2sc 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_107__tp4/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py index 9f79f4b69517..dbfb06ea7c2d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py @@ -12,7 +12,7 @@ class M3SelectLongDsl(OpWrapper): - """``m3_select_long``'s contract, cells and reference, run by ``kernels/m3_select_long_dsl``: the same selection + """``m3_select_long``'s contract, cells and reference, run by ``ops/m3_select_long_dsl``: the same selection bit for bit as the 1024-thread ``minimax_m3_sm107_tp4_select_long`` (also on rows where more than 2048 blocks tie at the 16th-largest score), from one 384-thread CTA per row. ``fast_tail`` (the 1024-thread kernel's warp-0 rank) does not apply and is ignored. Returns ``out``.""" @@ -102,7 +102,7 @@ def is_valid( The kernel is compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; the target compiles its keys before warmup - (`kernels/m3_select_long_dsl.prebuild`), and a call outside them compiles + (`ops/m3_select_long_dsl.prebuild`), and a call outside them compiles on first use. """ @@ -118,7 +118,7 @@ def __call__( trigger_phase: int = 0, fast_tail: bool = True, ) -> torch.Tensor: - from ...kernels.m3_select_long_dsl import select_long_dsl + from ..m3_select_long_dsl import select_long_dsl if scores.shape[1] > MAX_BLOCKS: raise RuntimeError( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py index 04d76f241b29..96171a6a9a29 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py @@ -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_107__tp4/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py index 208cb25aa72e..0b999befbadc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py @@ -135,7 +135,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 (modeling.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_107__tp4/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/embedding_firstnorm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp4/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/embedding_firstnorm/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py index 6fb8b1572d22..e4cd998d587d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py @@ -53,7 +53,7 @@ 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 (ops/m3_draft_lowm.prebuild). """ ARCHS = frozenset({Arch.SM_107}) @@ -95,13 +95,9 @@ def __call__( run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) elif own_splitk: - # This target's copy of the split-K kernel (kernels/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` # it also pushes the rows into a drafter boundary all-reduce (modeling.DRAFT_AR_PUSH). - from ...kernels.m3_lowm_splitk import ( - SplitKTactic, - run_splitk_dense, - run_splitk_dense_push, - ) + from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push if ar is not None: run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py index c6ccde3e5ae4..59bfa0780a00 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py @@ -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_107__tp4/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py index f8b2ab00217f..615a6b2e6785 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py @@ -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_107__tp4/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/long_select.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/long_select.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/long_select.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_attn_lc.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_attn_lc.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_attn_lc.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_lowm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py index 00c6c6887575..e89f2de8a8fb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py @@ -74,7 +74,7 @@ def draft_linear( weight: torch.Tensor, x: torch.Tensor, bias: Optional[torch.Tensor] ) -> Optional[torch.Tensor]: """x [M, K] @ weight^T through the low-M kernel when the target covers it (``prebuild``'s keys); None for the - plain GEMM. The split-K ('s') tactics run this target's copy of FlashInfer's kernel (kernels/m3_lowm_splitk, bit + plain GEMM. The split-K ('s') tactics run this target's copy of FlashInfer's kernel (ops/m3_lowm_splitk, bit for bit the same output at every certified cell), which also carries the drafter-boundary push epilogue.""" if bias is not None or x.dim() != 2: return None diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_draft_step.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_drafter_attn.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py index e5fa1bb5e945..7936d2497731 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_drafter_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py @@ -8,7 +8,7 @@ This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context -PrimTS split-KV decode of ``kernels/m3_fullctx_decode`` at the same pinned recipes as the dense layers. +PrimTS split-KV decode of ``ops/m3_fullctx_decode`` at the same pinned recipes as the dense layers. With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead (``minimax_m3_sm107tp4_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fold_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_force_vec.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_fullctx_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score_t2sc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score_t2sc.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score_t2sc.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score_t2sc.py index 9c0757ed166a..52a79fbffca9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_index_score_t2sc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score_t2sc.py @@ -4,7 +4,7 @@ # Vendored from vLLM (Apache-2.0): # https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py """MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule: T2SC, this target's copy of -the sm_107 TP2 / TP8 targets' scorer (``modeling.INDEX_SCORE_T2SC``). The sm_103-derived ``kernels/m3_index_score`` +the sm_107 TP2 / TP8 targets' scorer (``modeling.INDEX_SCORE_T2SC``). The sm_103-derived ``ops/m3_index_score`` stays the kill path. The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_lowm_splitk.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py index bf20343f7886..11ef3a496132 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_lowm_splitk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py @@ -3,7 +3,7 @@ """Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. This target's copy of FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the Eagle3 -drafter's low-M linears, ``kernels/m3_draft_lowm``): the same kernel and the same split order at every tactic, +drafter's low-M linears, ``ops/m3_draft_lowm``): the same kernel and the same split order at every tactic, bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's final cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, the SwiGLU prologue that follow. @@ -317,7 +317,7 @@ def _store_shared_remote_v4( # The drafter-boundary push (``ar_push``, modeling.DRAFT_AR_PUSH): the exchange words of -# kernels/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +# ops/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. @dsl_user_op def _ld_volatile_s32(addr, *, loc=None, ip=None): return Int32( @@ -935,7 +935,7 @@ def _push_tile( ar_n: cutlass.Int32, ): """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output - would store them (kernels/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + would store them (ops/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" cute_ext.partition_and_copy(thr_t2r, rD, sD) cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_mc_allreduce.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe16_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe16_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe16_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe16_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/m3_moe_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_moe/op16.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_prequant.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_report.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_report.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_report.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_sampled_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_sampled_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_sampled_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_select_long_dsl.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_shard_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_spec_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_target_head.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/m3_tokmap_guard.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py index 871f6887b7a5..258d45007a9b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py @@ -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_107__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py index f7b53b205fd9..7fe1bcb50976 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py @@ -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_107__tp4/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/msa_kv_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/m3_mx_gemv_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp4/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/mx_gemv/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py index 7ba04c92eb61..31c499cab2a6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py @@ -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_107__tp4/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py index fcb327e374e3..b980aeefc8d8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py @@ -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_107__tp4/ops/spec/m3_force_vec_req.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py index 5a2ffd8b6627..416425f9a01d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py @@ -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_req_kernel +from ..m3_force_vec import force_vec_req_kernel class M3ForceVecReq(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/kernels/triton_sparse_decode_mq.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/__init__.py deleted file mode 100644 index c5fd79e59d7f..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107 / tp8 target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py index 9a279b669def..490687087aee 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py @@ -37,12 +37,12 @@ range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise ``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). A row range is empty unless a landed cell has certified it at the TP8 geometry on Rubin. Every fast path a GEN rank takes prints a rank-tagged -engagement line once per (path, rows) (``taken``). Kernel modules report through ``kernels/m3_report``, which this +engagement line once per (path, rows) (``taken``). Kernel modules report through ``ops/m3_report``, which this module wires to the same policy at import. **Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the -long-context decode for the o_proj. It publishes it in ``kernels/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +long-context decode for the o_proj. It publishes it in ``ops/m3_prequant.SLOT`` keyed by the BF16 tensor; the next MXFP8 projection on a decode graph takes it, and every take clears the slot. The decode boundaries' push targets (the o_proj / dense-down / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's exchange) are passed explicitly from the layer to the producer and back. @@ -161,8 +161,8 @@ 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 ( +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import ( m3_drafter_attn, m3_fc_wait, m3_fold_tail, @@ -171,14 +171,12 @@ m3_report, m3_spec_tail, ) -from .kernels import m3_fullctx_decode as m3_fullctx_decode_kernel -from .kernels import m3_index_score as _index_score_kernel -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 import triton_sparse_decode as _sparse_decode_kernel -from .kernels.long_select import long_select_blocks -from .kernels.mx_gemv import op as mx_gemv +from .ops import m3_fullctx_decode as m3_fullctx_decode_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_shard_accept as m3_shard_accept_kernel +from .ops import msa_kv_update as _msa_kv_update_kernel +from .ops import triton_sparse_decode as _sparse_decode_kernel from .ops.activation.m3_silu_and_mul import m3_silu_and_mul from .ops.attention.m3_draft_attention import m3_draft_attention from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope @@ -217,8 +215,10 @@ from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.long_select import long_select_blocks from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm from .ops.spec.m3_capture_norm import m3_capture_norm from .ops.spec.m3_chunk_norm import m3_chunk_norm @@ -255,7 +255,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). # The multicast all-reduce + residual + RMSNorm at both boundaries of every TP2 decode step (one 8-token bucket), the # MoE boundary's also publishing the next QKV activation in MXFP8. SM107 TP2 hop: 3.5..4.4 us at 1..8 rows against # 8.7..9.5 us for the stock AUTO (NCCL) all-reduce + fused residual RMSNorm, and 13.3..14.0 us at the MoE boundary @@ -280,7 +280,7 @@ class DecodePlan: # TP2 skeleton: no cap (the built-in split) until a TP2 cap is measured and certified. DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) -# The dense layers' decode attention on PrimTS split-KV decode (kernels/m3_fullctx_decode) instead of the trtllm-gen +# The dense layers' decode attention on PrimTS split-KV decode (ops/m3_fullctx_decode) instead of the trtllm-gen # call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the # BF16 merge. At TP8 (8 query heads, 1 KV head per rank) the stock split is capped at 16 cluster CTAs per (request, KV # head), i.e. 16 of 212 CTAs at one request. Rubin TP8 geometry, FP8 KV, chained us per call at 165k / 323k, every @@ -331,7 +331,7 @@ class DecodePlan: # The K/V length bound the full-context plans are compiled for; a step whose cache length bound is longer is a miss. It # covers the TP2 AgentX config's max_seq_len (1 << 20) plus the speculative rows the cache reserves past it. FULLCTX_DECODE_MAX_KV_LEN = 1 << 21 # capacity: fixed: the full-context plans' compiled K/V bound -# The Eagle3 drafter layer's generation phase on the same full-context decode (kernels/m3_drafter_attn: the in-tree +# The Eagle3 drafter layer's generation phase on the same full-context decode (ops/m3_drafter_attn: the in-tree # RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its # split at 8 CTAs per (request, KV head) on SM 10.7: 45-46 us per draft pass at 165k. DRAFTER_FULLCTX_ENABLED = True @@ -344,7 +344,7 @@ class DecodePlan: # thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request # multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by # that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, -# token-0 drafts). kernels/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# token-0 drafts). ops/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such # calls: what the build writes (padded rows included), so bitwise the op with a fresh map. False:. PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" # How the rotation's products are summed, as the in-tree preprocess build sums them (0: two products then the add, @@ -372,7 +372,7 @@ class DecodePlan: # asserts it). PRODUCER_PREWAIT_STEP_INPUTS = True # #117: on a decode step of at most 8 rows the sparse layer's QKV + index GEMV runs the producer in its owner epilogue -# (kernels/mx_gemv m3_mx_gemv_producer, m3_layer.phase_producer.producer_head): one launch and one PDL edge fewer per +# (ops/mx_gemv m3_mx_gemv_producer, m3_layer.phase_producer.producer_head): one launch and one PDL edge fewer per # sparse layer, bit for bit the GEMV + m3_producer chain (q, index-q and every cache byte). The QKV linear defers its # GEMV (``_linear``'s ``defer``: a ``_DeferredGemv``) and the producer call site runs the fused kernel, or the plain # GEMV when it declines. @@ -424,13 +424,13 @@ class DecodePlan: # TP8 greedy early column, same-job A,B + B,A, flat (period +20..+26 at 155K..323K); True turns it on. MOE_EARLY_AR_RELEASE = False # A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel -# over the target lm_head's BF16 shard logits (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 # 0: (the gathered logits + argmax run). Rubin TP2 greedy verify (microbench): stock 120.4 us at b1; fused # with the coded head 117.0 us; the BF16-head + shard-accept arm is not measured yet. SHARD_ACCEPT_MAX_REQUESTS = 2 # capacity: fixed: the TP2 GEN batch # A sampled (temperature + top-p, advanced_sampling_mode no_topk) Eagle3 verify step of up to this many requests -# verifies with the fused sampled acceptance (kernels/m3_sampled_accept, the TP2-shard kernel) over the target +# verifies with the fused sampled acceptance (ops/m3_sampled_accept, the TP2-shard kernel) over the target # lm_head's BF16 shard logits, at the configured draft length (the AgentX contract). 0: (the gathered logits # and the in-tree sampler run). Rubin TP2 b1/b2 (p5 microbench): 180.1 / 181.8 -> 115.1 / 113.9 us per verify. SAMPLED_VERIFY_MAX_REQUESTS = 2 # capacity: fixed: the TP8 GEN batch @@ -450,7 +450,7 @@ class DecodePlan: # to this many verify rows; wider first passes take the in-tree first drafter inputs. 0:. CHUNK_NORM_MAX_ROWS = 8 # capacity: fixed: the TP2 verify rows (2 requests x 4) # The fused draft loop's two all-reduce boundaries per pass (o_proj -> post-attention norm, down -> final norm) of up -# to this many rows run the multicast all-reduce + residual + RMSNorm (kernels/m3_mc_allreduce, opts bit 16: the +# to this many rows run the multicast all-reduce + residual + RMSNorm (ops/m3_mc_allreduce, opts bit 16: the # drafter's stock rounding, normalizing the unrounded FP32 sum as FlashInfer fused_add_rmsnorm does) instead of the # built-in all-reduce + RMSNorm. Kill: 0. Rubin TP2 (p16 microbench): 10.0..11.0 -> 4.6..5.3 us per boundary. DRAFT_MC_AR_MAX_ROWS = 8 # capacity: fixed: the TP8 verify rows @@ -480,9 +480,9 @@ class DecodePlan: DI_FC_FLAG = True DI_FC_FLAG_PATH = "m3_di_fc_flag" # A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (spec_worker): per pass the -# draft input in one launch (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. @@ -501,13 +501,13 @@ class DecodePlan: # in-tree cuBLAS GEMM. Kill: 0. Rubin TP8 shards (sweep 777890, mclk locked) at 1 row: qkv 6.15 -> 3.94, # o 4.97 -> 2.42, gate_up 8.09 -> 5.78, down 5.62 -> 4.00 us; the replicated fc 21.50 -> 17.77 (TP2). DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the certified low-M rows -# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (kernels/m3_lowm_splitk): bit for bit the +# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (ops/m3_lowm_splitk): bit for bit the # same output at every certified cell (spec/test_modeling_v2_m3_lowm_splitk.py); the seam for the drafter-boundary push # epilogue and the down projection's SwiGLU prologue. Kill: False (FlashInfer's kernel). DRAFT_LOWM_OWN_SPLITK = True # On top of DRAFT_SILU_CPP: where the drafter MLP's down projection runs on a split-K tactic of the target's copy # (DRAFT_LOWM_OWN_SPLITK), its activation is built in the GEMM's B load from gate | up instead of a separate launch -# (kernels/m3_draft_lowm.draft_down_swiglu): bit for bit the activation + GEMM +# (ops/m3_draft_lowm.draft_down_swiglu): bit for bit the activation + GEMM # (spec/test_modeling_v2_m3_lowm_splitk.py), one launch fewer per draft pass. Kill: False. OFF in t8-3b: bitwise but # slower (SpecTail push-6b 782695, world 8, chained: 10.14 / 10.81 -> 17.38 / 66.85 us per boundary at 1 / 4 rows, # ~ +70 us/step); @@ -518,7 +518,7 @@ class DecodePlan: for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() } # The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows -# (kernels/triton_sparse_decode.MQ_MIN_TOKENS); the multi-query producer needs one KV head, which the TP8 heads have +# (ops/triton_sparse_decode.MQ_MIN_TOKENS); the multi-query producer needs one KV head, which the TP8 heads have # (8 query heads on 1 KV head per rank), so identity steps of at least 17 rows take it. SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover SPARSE_DECODE_KV_HEADS = 1 # capacity: fixed: the tp8 KV heads per rank @@ -551,7 +551,7 @@ class DecodePlan: # 16-CTA cluster launch and prologue then overlap the rank: Rubin TP2 in-model, locked memory clock, AgentX draws, # 2 instances vs the r4.3 stored A: replay period -162 us fixed, -94 at a 323K-token request. SELECT_LONG_TRIGGER_PHASE = 1 -# The long-row select runs as kernels/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per row +# The long-row select runs as ops/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per row # instead of one 1024-thread CTA, the same selection bit for bit, with 4, 11 or 24 keys per thread # (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices. On the Rubin TP8 decode chain (57 layers behind the # producer stream, 1 / 2 requests) the select is -0.15 / -0.12 us per layer at 165K tokens, 0.00 / -0.01 at 323K and @@ -582,10 +582,10 @@ class DecodePlan: # sequence). FORCE_AL_PER_REQUEST = True -# The decode MXFP8 GEMVs (kernels/mx_gemv, SM107 TP2 table), per projection (N, K): the rows each serves, on an +# The decode MXFP8 GEMVs (ops/mx_gemv, SM107 TP2 table), per projection (N, K): the rows each serves, on an # activation its producer quantized (the MoE boundary for the QKV, the attention merge for the o_proj) and on a bf16 # one (quantized in the GEMV's prologue, bit-identical to mxfp8_quantize). SM107 graph us per projection at 1..8 rows: -# 2.2..8.6 against 4.7..10.5 for the stock MXFP8 quantize + FlashInfer CuTe DSL GEMM (kernels/mx_gemv/op.py). +# 2.2..8.6 against 4.7..10.5 for the stock MXFP8 quantize + FlashInfer CuTe DSL GEMM (ops/mx_gemv/op.py). # False puts every projection back on the stock decode-graph GEMM. # On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the # built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. @@ -714,7 +714,7 @@ def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: m3_sparse_decode, _sparse_decode_producer, ) -# The TP8 long-context sparse decode at 8 query heads over one KV head (kernels/m3_attn_lc8): attention, split-K merge +# The TP8 long-context sparse decode at 8 query heads over one KV head (ops/m3_attn_lc8): attention, split-K merge # and the MXFP8 o_proj hand-off in one launch, for every decode step with a selected table; "prewait" up to # m3_attn_lc8.PREWAIT_MAX_ROWS query rows, "plain" above. Exact-class (in the float64 reference's band; run-to-run # deterministic). An empty range routes every step back to the Triton producer + merge. @@ -902,7 +902,7 @@ def _verify_rows(max_rows: int) -> FrozenSet[int]: ) # The dense layers' down projection takes the gate_up GEMV's output and applies the swigluoai activation in its -# quantize prologue (kernels/mx_gemv x_swiglu, bit for bit the Triton silu_and_mul kernel), instead of a separate +# quantize prologue (ops/mx_gemv x_swiglu, bit for bit the Triton silu_and_mul kernel), instead of a separate # silu_and_mul launch between the two GEMVs; rows the bf16 GEMV serves. Kill constant: False runs the activation # kernel as before. # OFF: bitwise but SLOWER (Host TP2 same-job pair 782860 / 782935: +60 / +65 us/step fixed, both orders; every GEMV CTA @@ -1105,7 +1105,7 @@ def reset_logged() -> None: _logged.clear() -# Kernel modules report from inside their calls through kernels/m3_report: the same policy. +# Kernel modules report from inside their calls through ops/m3_report: the same policy. m3_report.install( on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed ) @@ -1204,8 +1204,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. @@ -2543,9 +2543,9 @@ def prebuild_kernels(self) -> Dict[str, object]: ) if not self.mxfp8.quant_method.use_cutlass: raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") - from .kernels import m3_fullctx_decode as _fullctx_decode_kernel - from .kernels import m3_select_long_dsl as _select_long_dsl_kernel - from .kernels import m3_tokmap_guard + from .ops import m3_fullctx_decode as _fullctx_decode_kernel + from .ops import m3_select_long_dsl as _select_long_dsl_kernel + from .ops import m3_tokmap_guard device = self.w["embed"].device # The one-request token-map guard of the stock generation preprocess (``PREPROCESS_TOKMAP_FIX``): every path @@ -3750,7 +3750,7 @@ def _dense_attention(self, i: int, position_ids, hidden, step, push): def _fullctx_plan(self, i: int, device: torch.device, step): """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs - on ``kernels/m3_fullctx_decode`` (``FULLCTX_DECODE``): a covered (query rows per request, requests) shape over + on ``ops/m3_fullctx_decode`` (``FULLCTX_DECODE``): a covered (query rows per request, requests) shape over the layer's FP8 pool. None otherwise (the MSA backend runs). Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 (29.3), (1, 1) @@ -3775,7 +3775,7 @@ def _fullctx_plan(self, i: int, device: torch.device, step): if pool.dtype != torch.float8_e4m3fn: missed(path, num_tokens, "the dense layer's KV pool is not FP8") return None - from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + from .ops import m3_fullctx_decode as _fullctx_decode_kernel own_reducer = FULLCTX_DECODE_OWN_REDUCER if own_reducer and not _fullctx_decode_kernel.has_own_reducer(): @@ -3797,7 +3797,7 @@ def _fullctx_decode(self, i: int, q, k, v, step, plan_tokens, *, kv_written: boo write_msa_phase_kv, ) - from .kernels import m3_fullctx_decode as _fullctx_decode_kernel + from .ops import m3_fullctx_decode as _fullctx_decode_kernel core = self.core plan, num_tokens = plan_tokens @@ -3980,7 +3980,7 @@ def _select(self, i: int, max_score, step, num_tokens: int): if group > 1 else max_score ) - from .kernels import m3_select_long_dsl as _select_long_dsl_kernel + from .ops import m3_select_long_dsl as _select_long_dsl_kernel num_blocks = int(max_score_kv.shape[1]) table = None @@ -4032,11 +4032,11 @@ def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise the index scorer writes the per-index-head block maxima, reduced to the KV head, and the long-row selector takes the top 16 (the built-in selector beyond its shapes). The long-context decode (``m3_sparse_decode`` on - ``kernels/m3_attn_lc8``) prefetches into L2 this layer's o_proj weights (``SPARSE_DECODE_LC8_PREFETCH_OPROJ``) + ``ops/m3_attn_lc8``) prefetches into L2 this layer's o_proj weights (``SPARSE_DECODE_LC8_PREFETCH_OPROJ``) and the next sparse layer's K/V pages of the blocks that layer selected one step back (``SPARSE_DECODE_LC8_PREFETCH_NEXT_KV``, ``_next_kv_prefetch``); hints only. """ - from .kernels import m3_attn_lc8 as _attn_lc8_kernel + from .ops import m3_attn_lc8 as _attn_lc8_kernel core = self.core backend = core.attn[i] @@ -4363,7 +4363,7 @@ class M3Eagle3Draft(nn.Module): (``ops/gemm/m3_draft_lowm``, the TP8 port) at the rows ``DRAFT_LOWM`` covers, compiled before the warmup (``prebuild_linears``); every other call is the in-tree cuBLAS GEMM. Rubin TP2 (mclk locked, 4 rotating sets) at 1 row: qkv 13.18 -> 9.71, o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. The - split-K ('s') tactics run this target's copy of FlashInfer's kernel (``kernels/m3_lowm_splitk``, bit for bit the + split-K ('s') tactics run this target's copy of FlashInfer's kernel (``ops/m3_lowm_splitk``, bit for bit the same output); with ``DRAFT_AR_PUSH`` the o_proj and down projection of the fused draft loop also push their rows into the drafter boundary all-reduce's exchange (``arm_push`` / ``take_pushed``); the direct ('d') tactics never push. @@ -4534,7 +4534,7 @@ def _lowm_linear(self, weight: torch.Tensor, x: torch.Tensor) -> Optional[torch. def prebuild_linears(self) -> int: """Compile (and run once) every covered row count of every drafter projection, outside capture, on the current stream, before the executor's warmup. Returns the number of calls compiled.""" - from .kernels.m3_lowm_splitk import compile_push + from .ops.m3_lowm_splitk import compile_push use_pdl = get_env_enable_pdl() built = 0 @@ -4608,7 +4608,7 @@ def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is False, where the caller fuses the reduction). The backend's generation phase runs this target's full-context - decode with the fused RoPE + FP8 K/V append when the shell installed it (``kernels/m3_drafter_attn``).""" + decode with the fused RoPE + FP8 K/V append when the shell installed it (``ops/m3_drafter_attn``).""" qkv = self.linear(x, "qkv") _record("m3_draft_attention", int(md.num_tokens)) o = m3_draft_attention(self.attn, qkv, md) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py index a5439efc660b..88ce9728f40b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The dense layers' full-context decode attention at TP8 on Rubin (``kernels/m3_fullctx_decode``): PrimTS split-KV +"""The dense layers' full-context decode attention at TP8 on Rubin (``ops/m3_fullctx_decode``): PrimTS split-KV decode at the recipe modeling's ``FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" from __future__ import annotations @@ -12,7 +12,7 @@ fp8_e4m3_saturate, ) -from ...kernels import m3_fullctx_decode as kernel +from .. import m3_fullctx_decode as kernel from .m3_dense_decode import dense_attention_reference, gather_paged_kv __all__ = ["m3_fullctx_decode"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py index 674b16e7744f..5eb42546612e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py @@ -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 @@ -26,7 +26,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_107__tp8/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py index cd76a4e6d5d2..7d3627bba5f3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py @@ -12,7 +12,7 @@ class M3SelectLongDsl(OpWrapper): - """``m3_select_long``'s contract, cells and reference, run by ``kernels/m3_select_long_dsl``: the same selection + """``m3_select_long``'s contract, cells and reference, run by ``ops/m3_select_long_dsl``: the same selection bit for bit as the 1024-thread ``minimax_m3_sm107_tp2_select_long`` (also on rows where more than 2048 blocks tie at the 16th-largest score), from one 384-thread CTA per row. Returns ``out``.""" @@ -105,7 +105,7 @@ def is_valid( The kernel is compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; the target compiles its keys before warmup - (`kernels/m3_select_long_dsl.prebuild`), and a call outside them compiles + (`ops/m3_select_long_dsl.prebuild`), and a call outside them compiles on first use. """ @@ -120,7 +120,7 @@ def __call__( early_trigger: bool = False, trigger_phase: int = 0, ) -> torch.Tensor: - from ...kernels.m3_select_long_dsl import select_long_dsl + from ..m3_select_long_dsl import select_long_dsl if scores.shape[1] > MAX_BLOCKS: raise RuntimeError( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py index 6525a3341b15..fc770d010643 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py @@ -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_107__tp8/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py index 68773b5b74f2..cf0a4be970d9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py @@ -97,7 +97,7 @@ class M3McAllreduceNorm(OpWrapper): ARCHS = frozenset({Arch.SM_107}) # The tp2 decode boundaries (hidden 6144, 2 ranks, eps 1e-6): every decode step fits the 8-token bucket, epilogue - # option 7 at 1, 2, 7, 8 tokens and 15 at 3..6 (kernels/m3_mc_allreduce.AR_OPTS_BY_M; every option is + # option 7 at 1, 2, 7, 8 tokens and 15 at 3..6 (ops/m3_mc_allreduce.AR_OPTS_BY_M; every option is # bit-identical). CELLS: tuple[Cell, ...] = ( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/embedding_firstnorm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp8/kernels/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/embedding_firstnorm/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py index 7e0e023ee88c..8d4a56e3b437 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py @@ -97,13 +97,9 @@ def __call__( run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) elif own_splitk: - # This target's copy of the split-K kernel (kernels/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` # it also pushes the rows into a drafter boundary all-reduce (modeling's DRAFT_AR_PUSH). - from ...kernels.m3_lowm_splitk import ( - SplitKTactic, - run_splitk_dense, - run_splitk_dense_push, - ) + from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push if ar is not None: run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py index 4e9f8f65de1b..a16702f2cc1a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py @@ -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_107__tp8/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py index 3b5d21235f29..159b49ff69af 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py @@ -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_107__tp2/kernels/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/long_select.py similarity index 98% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/long_select.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/long_select.py index 14e24fbdbcac..79aaf6674300 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/kernels/long_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/long_select.py @@ -14,7 +14,7 @@ The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long`` (this target's copy of ``minimax_m3_select_long`` with a choice of where it releases the sparse decode), called through this target's ``m3_select_long`` entry, or with -``dsl`` the same selection from one 384-thread CTA per row (``m3_select_long_dsl``, kernels/m3_select_long_dsl). +``dsl`` the same selection from one 384-thread CTA per row (``m3_select_long_dsl``, ops/m3_select_long_dsl). """ from typing import Optional diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_attn_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_attn_lc8.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_attn_lc8.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_attn_lc8.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_draft_step.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_drafter_attn.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py index 782ff2afb816..e582e2309e61 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_drafter_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py @@ -8,7 +8,7 @@ This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context -PrimTS split-KV decode of ``kernels/m3_fullctx_decode`` at the same pinned recipes as the dense layers. +PrimTS split-KV decode of ``ops/m3_fullctx_decode`` at the same pinned recipes as the dense layers. With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead (``minimax_m3_sm107tp8_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fc_wait.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fc_wait.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fc_wait.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fc_wait.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fold_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_force_vec.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_fullctx_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_index_score.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/phase_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_layer/phase_producer.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_locality.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_lowm_splitk.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py index 685797a96700..2f1a599c45aa 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_lowm_splitk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py @@ -390,7 +390,7 @@ def _fence_proxy_async_global(*, loc=None, ip=None) -> None: # The drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH): the exchange words of -# kernels/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +# ops/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. @dsl_user_op def _ld_volatile_s32(addr, *, loc=None, ip=None): return Int32( @@ -1103,7 +1103,7 @@ def _push_tile( ar_n: cutlass.Int32, ): """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output - would store them (kernels/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + would store them (ops/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" cute_ext.partition_and_copy(thr_t2r, rD, sD) cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_mc_allreduce.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/m3_moe_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_moe/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_prequant.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_report.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_report.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_report.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_router_gemv.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_sampled_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_sampled_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_sampled_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_select_long_dsl.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_shard_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_spec_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_spec_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_spec_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_target_head.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/m3_tokmap_guard.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py index 88eb4f183b3a..113878a991f9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py @@ -17,7 +17,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_107__tp8/kernels/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/msa_kv_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/m3_mx_gemv_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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_107__tp8/kernels/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/mx_gemv/op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py index c2e14b5d8c4b..11b92b47fc50 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py @@ -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_107__tp8/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py index 2001bf183b54..a192e89c5c76 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__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, force_vec_req_kernel +from ..m3_force_vec import force_vec_kernel, force_vec_req_kernel class M3ForceVec(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/kernels/triton_sparse_decode_mq.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py index 1bd4e703e9dc..2dfd7aa71ebb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """The MSA page table is written on the host side only, before the graph replay. -The index scorer's early issue (``kernels/m3_index_score.py``, ``early_issue``) reads ``block_table[b, split]`` BEFORE +The index scorer's early issue (``ops/m3_index_score.py``, ``early_issue``) reads ``block_table[b, split]`` BEFORE its PDL grid-dependency wait and loads that page's K tile right after it. That is exact only while no kernel of the step writes ``MsaMetadata.msa_block_table``: today its one writer is the host-to-device copy in ``_build_msa_fields`` (called from ``prepare()``, before the replay), besides the allocation. This test scans @@ -99,7 +99,7 @@ def test_only_prepare_writes_the_msa_page_table() -> None: for func, line in _write_sites(ast.parse(src)): assert (rel, func) in _ALLOWED, ( f"{rel}:{line} ({func}) writes msa_block_table; the index scorer's early issue reads it before its " - "PDL grid wait (kernels/m3_index_score.py), so an in-graph writer must be ordered before it or the " + "PDL grid wait (ops/m3_index_score.py), so an in-graph writer must be ordered before it or the " "early page read moved after the wait" ) found.add((rel, func)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py index 069e4fdc3063..6a42e13000de 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test (SM 10.7) for kernels/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +"""GPU test (SM 10.7) for ops/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row first draft pass and the 1-row later passes of 1-2 requests. @@ -20,10 +20,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_drafter_attn as drafter, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_fullctx_decode as fullctx, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py index a27ede466e03..3d77b932b3e7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -16,10 +16,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_drafter_attn as drafter, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_fullctx_decode as fullctx, ) from tensorrt_llm._torch.attention.backends.fmha.utils import ( @@ -376,10 +376,10 @@ def test_crafted_inputs_pin_the_rope_mode() -> None: def test_one_request_multi_token_matches_the_guarded_preprocess() -> None: """The drafter's first pass at one request (several tokens): the in-tree preprocess places those rows by a - workspace token map it does not build for one request, which the TP2 token-map guard (kernels/m3_tokmap_guard) + workspace token map it does not build for one request, which the TP2 token-map guard (ops/m3_tokmap_guard) writes right before the call. After POISONING the map with a two-request call, the fused op must be bitwise the guarded preprocess + cast; this is what drafter.ROPE_ONE_REQUEST_MULTI rests on.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_tokmap_guard as guard, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py index cb20a22a51d0..6a21693bcf9a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py @@ -112,7 +112,7 @@ def test_the_cells_cover_the_target_rows() -> None: def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) @@ -152,7 +152,7 @@ def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode step first hit (NaN rows from exp2(-inf - -inf)).""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl @@ -186,7 +186,7 @@ def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py index baf8ef7a2b73..b883a15acd38 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """The per-token valid-block count is written only before the step's first layer. -The long-row select (``kernels/m3_select_long_dsl.py``, ``SELECT_N_VALID_PREWAIT``) reads its row's entry of +The long-row select (``ops/m3_select_long_dsl.py``, ``SELECT_N_VALID_PREWAIT``) reads its row's entry of ``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the buffer happens before the step's first layer: @@ -131,7 +131,7 @@ def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: for func, line in _write_sites(ast.parse(src)): assert (rel, func) in _ALLOWED, ( f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " - "wait (kernels/m3_select_long_dsl.py, SELECT_N_VALID_PREWAIT), so a writer inside the " + "wait (ops/m3_select_long_dsl.py, SELECT_N_VALID_PREWAIT), so a writer inside the " "step must be ordered before the select's grid or the read moved after the wait" ) found.add((rel, func)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py index d452cb59f998..60b04b9dde33 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( m3_select_long_dsl as select_kernels, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long import ( # noqa: E501 @@ -140,7 +140,7 @@ def test_trigger_phase_selects_the_same(cell: int) -> None: @pytest.mark.parametrize("early_trigger", [False, True]) @pytest.mark.parametrize("cell", range(len(op.CELLS))) def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: - """kernels/m3_select_long_dsl.SELECT_N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the + """ops/m3_select_long_dsl.SELECT_N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the grid wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and tied rows).""" phase = 0 if early_trigger else 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py index 3fd3f605a230..feb450a5d350 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py @@ -17,11 +17,11 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( modeling, # noqa: E501 ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_prequant, triton_sparse_decode, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 triton_sparse_decode as tsd, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py index b726d37e1dd3..477b7ef94ebf 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py @@ -319,7 +319,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def _target_decode(c, killed): - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_prequant, triton_sparse_decode, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py index 6d4616218e0b..3a19b3bf5a7d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test (SM 10.7) for kernels/m3_tokmap_guard: a one-request multi-token call of +"""GPU test (SM 10.7) for ops/m3_tokmap_guard: a one-request multi-token call of thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for @@ -16,7 +16,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_tokmap_guard as guard, ) from tensorrt_llm.bindings.internal import thop diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py index cf8c5ab15574..d35556356264 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py @@ -7,15 +7,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.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_107__tp2.ops.attention.msa_kv_update import ( # noqa: E501 TOKENS_PER_BLOCK, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.msa_kv_update import ( # noqa: E501 msa_kv_update as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py index 2e09105e8fcb..8e0e595bdc1c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -329,7 +329,7 @@ def _capture_reference( updated: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor ) -> None: """The capture the boundary replaces: m3_spec_tail.capture_norm of the updated residual (no second addend).""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_spec_tail, ) @@ -644,7 +644,7 @@ def _run_one_rank() -> int: RANK, WORLD = rank, world assert WORLD == 2, f"the entries are certified at 2 ranks, got {WORLD}" - import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_mc_allreduce as _ar # noqa: E501 from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm import ( # noqa: E501 m3_mc_allreduce_norm as _plain, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py index 99df8fb953ef..c6baaa215341 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py @@ -298,8 +298,8 @@ def _run_one_rank() -> int: RANK, WORLD = rank, world assert WORLD == 2, f"the push is certified at 2 ranks, got {WORLD}" - import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_mc_allreduce as _ar # noqa: E501 - import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_prequant as _pq # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_prequant as _pq # noqa: E501 from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 sf_written as _sf_written, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py index 3af452daa1d9..f1d036a31909 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py @@ -14,13 +14,13 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.mx_gemv import ( # noqa: E501 - op as gemv, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 DENSE_DOWN, DENSE_GATE_UP, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.mx_gemv import ( # noqa: E501 + op as gemv, +) assert torch.cuda.is_available(), "the dense seam test requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py index a514ebbe4b38..46548f91c1d7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py @@ -17,7 +17,7 @@ def _locality(): - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_locality import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_locality import ( # noqa: E501 LocalityMap, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py index e7a31d3ae752..a3509b9ab5a9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Rule 1b gate a for the TP2 decode GEMV order override (kernels/mx_gemv/op.STRIDED_K_TILES: strided k-tiles): +"""Rule 1b gate a for the TP2 decode GEMV order override (ops/mx_gemv/op.STRIDED_K_TILES: strided k-tiles): for every overridden (projection, rows) the output is the same on 3 eager calls and on a CUDA graph replay (deterministic), and its max abs / max rel error vs the fp64 product of the dequantized operands stays within the shipped order's band (the shared table's config on the same operands, over several seeds).""" @@ -8,9 +8,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.mx_gemv import ( # noqa: E501 - op, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 DENSE_GATE_UP, O_PROJ, @@ -21,6 +18,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 mxfp8_dequant, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.mx_gemv import ( # noqa: E501 + op, +) pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py index aad0252e6662..1d6de982a67a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py @@ -65,7 +65,7 @@ def _state(m: int) -> M3MoeState: def test_a_module_built_with_other_trace_config_compiles_its_own_kernel(monkeypatch) -> None: """The compile cache keys on the config the kernel module was built with: a state whose module carries an extra trace-time option (``timing``, which only adds per-CTA stamps) compiles its own kernel, never the plain build's.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_moe import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_moe import ( # noqa: E501 op as kop, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py index eb0f3beec25c..700ce9af2be9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Who finalizes the wide target's MoE at each generic row count: past the decode rows, up to GENERIC_MC_FINALIZE_MAX_TOKENS, the finalize (+ the shared expert) folds into the boundary's multicast all-reduce -(M3DecoderLayer._mc_finalize_boundary); from MOE_TRITON_FINALIZE_MIN_TOKENS, kernels/m3_moe_finalize finalizes it. +(M3DecoderLayer._mc_finalize_boundary); from MOE_TRITON_FINALIZE_MIN_TOKENS, ops/m3_moe_finalize finalizes it. The two ranges meet with no gap and no overlap, and the fold declines from shapes alone past its cap. CPU only: nothing here launches a kernel. @@ -16,7 +16,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_mc_allreduce, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py index cd34e710ecfd..bc95a2586d74 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py @@ -11,7 +11,7 @@ import pytest -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_moe import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_moe import ( # noqa: E501 op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py index 13cec78a3a07..0530ad9cc24e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py @@ -12,7 +12,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_moe import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_moe import ( # noqa: E501 op as engine_op, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py index cc92290a9162..5523a391661e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""kernels/m3_moe_finalize: the routed experts' finalize (+ the shared expert) from the TRTLLM-Gen runner's unfinalized +"""ops/m3_moe_finalize: the routed experts' finalize (+ the shared expert) from the TRTLLM-Gen runner's unfinalized parts, bitwise to the runner's finalizeKernelVecLoad order (acc = 0; acc += w_k * fc2[idx_k] for k in order, skipping -1; one BF16 rounding) and to ``shared.add_(routed)``, at the rows the wide target routes there (49..16K), with dropped (-1) slots, padded FC2 rows, eager and CUDA-graph replay. (The runner itself was compared at 512 / 2K / 4K / 16K rows on @@ -9,7 +9,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_moe_finalize import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_moe_finalize import ( # noqa: E501 m3_moe_finalize, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py index fde15dd95781..6b2c334a7615 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py @@ -4,13 +4,13 @@ The draft-step and sampled-accept workspaces are created collectively (both ranks, a multicast buffer) on the first eager call. Creating one inside a capture would allocate and synchronize mid-capture, so a capture that finds it missing -reports a miss (``kernels/m3_report.capture_missed``) and gets None (the caller then takes its stock path) instead. +reports a miss (``ops/m3_report.capture_missed``) and gets None (the caller then takes its stock path) instead. """ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( m3_draft_step, m3_report, m3_sampled_accept, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py index 30931023b5ed..deeb0b10c5d3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test of the drafter boundary push (modeling.DRAFT_AR_PUSH / DRAFT_AR_PUSH_DIRECT) on one rank: the push variants -of the split-K copy (kernels/m3_lowm_splitk) and the direct copy (kernels/m3_lowm_direct) write FlashInfer's output +of the split-K copy (ops/m3_lowm_splitk) and the direct copy (ops/m3_lowm_direct) write FlashInfer's output bit for bit, and store exactly the boundary all-reduce's canonical exchange words of that output: word (t, column pair) = bf16 lo | bf16 hi << 16 with -0.0 halves as +0.0, at parity * maxt * world * H / 2 + (t * world + rank) * H / 2 + column / 2, the parity being the (t, column / (H / 8)) flag's low bit; every other word stays untouched. Eager and in a @@ -15,7 +15,7 @@ from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiSplitK from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_splitk -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_lowm_direct, m3_lowm_splitk, ) @@ -28,7 +28,7 @@ H, MAXT, EMPTY = 6144, 8, -2147483648 # The boundary linears (output width H): o_proj and down, each certified row count on split-K or an even direct width, -# up to the 8-token bucket the push arms on (kernels/m3_draft_lowm.arm_push); the batch-4/8 rows past it run unpushed. +# up to the 8-token bucket the push arms on (ops/m3_draft_lowm.arm_push); the batch-4/8 rows past it run unpushed. CELLS = [ (nk[1], m, *tactic_for(nk[0], nk[1], m)) for nk, rows in sorted(ROWS.items()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py index 9669fc829def..964e62f86558 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py @@ -6,7 +6,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_input import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py index 7d0cc68a5011..0528b9f1abb3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py @@ -18,7 +18,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_draft_step as draft_step_kernels, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_target_head import ( # noqa: E501 @@ -687,7 +687,7 @@ def test_guard_refuses_too_few_ctas() -> None: op(*args) -# ------------------------------------------------- one pass as two launches (uGPU, kernels/m3_draft_step.UGPU) +# ------------------------------------------------- one pass as two launches (uGPU, ops/m3_draft_step.UGPU) @pytest.mark.parametrize( @@ -699,7 +699,7 @@ def test_two_launches_per_domain_are_the_one_launch_pass(cell) -> None: """A pass split into two launches (tiles [0, split) and [split, n), one per locality domain stream, the exchange and bookkeeping by the last of both launches' CTAs) writes exactly what the one-launch pass writes: tokens, winning keys, draft columns, next positions, KV and sequence lengths, and re-arms the counters and keys.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_locality import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_locality import ( # noqa: E501 LocalityMap, ) @@ -736,7 +736,7 @@ def test_two_launches_per_domain_are_the_one_launch_pass(cell) -> None: assert two[6].tolist() == [0, 1] and not bool(two[5].any()), "counters / keys not re-armed" -# ----------------------------------------------- equal-chunk split (kernels/m3_draft_step.EQ_CHUNKS: bnd / bnd_cnt set) +# ----------------------------------------------- equal-chunk split (ops/m3_draft_step.EQ_CHUNKS: bnd / bnd_cnt set) def _eq_buffers(n_tiles: int): @@ -757,7 +757,7 @@ def test_equal_chunk_split_is_the_whole_tile_pass(cell, launches) -> None: bnd: the same tokens, winning global ids, draft columns, positions and lengths as whole tiles per CTA (the logits differ at most by the order of fp32 additions: their BF16 keys by at most one ulp), and bnd / bnd_cnt are left zero for the next call, over one launch and over the two per-domain launches.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_locality import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_locality import ( # noqa: E501 LocalityMap, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py index 3c71a4a39f07..7c9a826c1d3d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the forward-tail spec fold (kernels/m3_fold_tail and M3SpecStepFold.in_forward). +"""GPU test for the forward-tail spec fold (ops/m3_fold_tail and M3SpecStepFold.in_forward). The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers (step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing @@ -15,7 +15,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( m3_fold_tail, m3_spec_tail, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py index e68903313514..6d585d8afa32 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (kernels/m3_lowm_splitk, +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (ops/m3_lowm_splitk, modeling.DRAFT_LOWM_OWN_SPLITK): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's linears, eager and in a CUDA graph, with and without PDL.""" @@ -9,14 +9,14 @@ from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels.m3_lowm_splitk import ( # noqa: E501 - SplitKTactic, - run_splitk_dense, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 ROWS, tactic_for, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, +) assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py index 2af99c2062f6..7a59742e1302 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py @@ -115,7 +115,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res @pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) def test_prewait_parameter_loads_change_nothing(cell) -> None: - """kernels/m3_sampled_accept.PREWAIT_PARAMS moves only where the sampling parameters load (before or after the PDL + """ops/m3_sampled_accept.PREWAIT_PARAMS moves only where the sampling parameters load (before or after the PDL wait): the same draws and acceptance outputs, bit for bit, call after call.""" outs = [] for prewait in (True, False): @@ -134,7 +134,7 @@ def test_prewait_parameter_loads_change_nothing(cell) -> None: @pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) def test_the_early_trigger_changes_nothing(cell) -> None: - """kernels/m3_sampled_accept.EARLY_TRIGGER moves only where the grid releases its dependents: the same draws and + """ops/m3_sampled_accept.EARLY_TRIGGER moves only where the grid releases its dependents: the same draws and acceptance outputs, bit for bit, call after call, PDL launch on.""" outs = [] for early in (True, False): @@ -167,7 +167,7 @@ def _walk_outputs(args, ilp, calls=3, offset_step=1000): @pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) def test_the_ilp_walk_changes_nothing(cell) -> None: - """kernels/m3_sampled_accept.ILP_WALK (SA-PW) only computes the owner's token weights 8 at a time; the walk adds + """ops/m3_sampled_accept.ILP_WALK (SA-PW) only computes the owner's token weights 8 at a time; the walk adds them in the same order with the same early exit, so the draws and acceptance outputs are the one-token walk's, bit for bit, call after call.""" one = _walk_outputs(_case(cell.spec, seed_of(cell.why)), False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py index c9cd78d40985..49e21a3a471b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py @@ -16,10 +16,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_draft_step as draft_step_kernel, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 triton_sparse_decode as tsd, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py index dfb89711db4d..37b04013094a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py @@ -14,7 +14,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_mc_allreduce as mc, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py index dcd2ffb2c38b..828bdba6d48a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py @@ -14,7 +14,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( triton_sparse_decode as tsd, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py index 6e22a6c33f36..e5cec6259f9a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test (SM 10.7) for kernels/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +"""GPU test (SM 10.7) for ops/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row first draft pass and the 1-row later passes of one request. @@ -20,13 +20,13 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_drafter_attn as drafter, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_fullctx_decode as fullctx, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_tokmap_guard as tokmap_guard, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py index 56d0a258d3ba..aefbea361b74 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -16,10 +16,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_drafter_attn as drafter, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_fullctx_decode as fullctx, ) from tensorrt_llm._torch.attention.backends.fmha.utils import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py index 6f2ace58365a..0014a17a7b01 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py @@ -112,7 +112,7 @@ def test_certified_cells(cell, subpages, certifying) -> None: # noqa: F811 def test_the_kv_stage_depth_keeps_the_bits(cell) -> None: """kv_stages only sets how many K/V tiles the load pipeline keeps in flight: the recipe's plan and plans compiled at 12 and at 4 stages give the same bits.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl @@ -157,7 +157,7 @@ def test_the_cells_cover_the_dispatch_rows() -> None: def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) @@ -197,7 +197,7 @@ def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode step first hit (NaN rows from exp2(-inf - -inf)).""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl @@ -231,7 +231,7 @@ def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py index 3c5b75f264ed..be10d54aca2a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """The per-token valid-block count is written only before the step's first layer. -The long-row select (``kernels/m3_select_long_dsl.py``, ``N_VALID_PREWAIT``) reads its row's entry of +The long-row select (``ops/m3_select_long_dsl.py``, ``N_VALID_PREWAIT``) reads its row's entry of ``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the buffer happens before the step's first layer: @@ -131,7 +131,7 @@ def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: for func, line in _write_sites(ast.parse(src)): assert (rel, func) in _ALLOWED, ( f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " - "wait (kernels/m3_select_long_dsl.py, N_VALID_PREWAIT), so a writer inside the " + "wait (ops/m3_select_long_dsl.py, N_VALID_PREWAIT), so a writer inside the " "step must be ordered before the select's grid or the read moved after the wait" ) found.add((rel, func)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py index 9866b1d589c0..816351089a99 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( m3_select_long_dsl as select_kernels, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long import ( # noqa: E501 @@ -140,7 +140,7 @@ def test_trigger_phase_selects_the_same(cell: int) -> None: @pytest.mark.parametrize("early_trigger", [False, True]) @pytest.mark.parametrize("cell", range(len(op.CELLS))) def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: - """kernels/m3_select_long_dsl.N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the grid + """ops/m3_select_long_dsl.N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the grid wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and tied rows).""" phase = 0 if early_trigger else 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py index cfd8a14643f6..1c72256dac47 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_report, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode import ( # noqa: E501 @@ -223,7 +223,7 @@ def test_identity_route_is_bitwise_the_stock_decode(requests, monkeypatch) -> No def _triton_route(monkeypatch) -> None: """Route selected steps to the target's Triton producer + merge (this target takes the LC whenever it applies).""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( m3_attn_lc, ) @@ -379,7 +379,7 @@ def test_lc_identity_route_reads_its_own_lengths_not_the_table( def test_lc_identity_step_reads_q_after_the_wait(requests, monkeypatch) -> None: """An identity step has no select, so the LC must read q after its grid wait: the op gets identity_step (and the device identity selection) on every identity step.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_attn_lc, m3_prequant, ) @@ -410,7 +410,7 @@ def test_lc_fails_closed_on_ids_past_the_row(requests) -> None: """A selected id past the token's row (a stale table, e.g. a padded row keeping a 500K request's n_valid) is skipped, as the Triton producer skips it: the rows equal those with that slot at -1, bit for bit, and stay finite even though the block-table tail it would reach points at a page whose block scales are NaN (0xFF).""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_prequant, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py index 35f04a37eb8e..013e12cd3703 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test (SM 10.7) for kernels/m3_tokmap_guard: a one-request multi-token call of +"""GPU test (SM 10.7) for ops/m3_tokmap_guard: a one-request multi-token call of thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for @@ -16,7 +16,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_tokmap_guard as guard, ) from tensorrt_llm.bindings.internal import thop diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py index 85bb1fd0bcde..a1682f829082 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py @@ -7,15 +7,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__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_107__tp4.ops.attention.msa_kv_update import ( # noqa: E501 TOKENS_PER_BLOCK, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.msa_kv_update import ( # noqa: E501 msa_kv_update as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py index 80fa9787751d..db9dba58f108 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py @@ -122,7 +122,7 @@ def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: def _capture_reference(updated, weight, eps, out) -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_spec_tail, ) @@ -217,7 +217,7 @@ def _run_one_rank() -> int: RANK, WORLD = rank, world assert WORLD == 4, f"the tp4 boundary runs at 4 ranks, got {WORLD}" try: - import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_mc_allreduce as _ar # noqa: E501 from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm import ( # noqa: E501 m3_mc_allreduce_norm as _plain, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py index 7e241a8c484e..13f7835c7d66 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py @@ -444,8 +444,8 @@ def _run_one_rank() -> int: RANK, WORLD = rank, world assert WORLD == 4, f"the drafter runs at tp4, got {WORLD} ranks" try: - import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_mc_allreduce as _ar # noqa: E501 import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.modeling as _modeling # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_mc_allreduce as _ar # noqa: E501 modeling, ar_module = _modeling, _ar assert modeling.m3_mc_allreduce is ar_module, ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py index 0b372cc02921..852149f12cc7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -554,7 +554,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_107__tp4.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_mc_allreduce as _ar # noqa: E501 from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm import ( # noqa: E501 m3_mc_allreduce_norm as _plain, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py index ddf72d7556e9..ebc36807e49b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py @@ -4,7 +4,7 @@ kernels (direct SIMT and split-K tensor-core), at the TP4 drafter shards, at the rows its table certifies. Every cell is held to the fp64 reference within the derived band. The drafter's projections (``M3Eagle3Draft._linear`` -over ``kernels/m3_draft_lowm``) are held on the drafter's weight table: a covered row count is the entry's call bit for +over ``ops/m3_draft_lowm``) are held on the drafter's weight table: a covered row count is the entry's call bit for bit, an uncovered one the plain GEMM's. """ @@ -17,7 +17,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_draft_lowm, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py index 03197b28202b..0cd67d512573 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py @@ -143,7 +143,7 @@ def test_prebuild_covers_the_split8_config_while_frozen(nk) -> None: from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( modeling as dispatch, ) - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.mx_gemv import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.mx_gemv import ( # noqa: E501 op as mx_gemv, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py index c1c7aad25bce..467cfd888e5e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py @@ -15,7 +15,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_target_head as target_head_kernels, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py index 2a82033abc0b..d999e3e1e83e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py @@ -12,7 +12,7 @@ import pytest -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_moe import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_moe import ( # noqa: E501 op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py index 47d3707f7182..2f82a2a9d63b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py @@ -6,7 +6,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_input import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py index 24987c779549..53aaf105b5a3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py @@ -1,13 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The drafter boundary push bookkeeping (kernels/m3_draft_lowm arm_push / take_pushed, dispatch.DRAFT_AR_PUSH): the +"""The drafter boundary push bookkeeping (ops/m3_draft_lowm arm_push / take_pushed, dispatch.DRAFT_AR_PUSH): the boundary all-reduce runs prepushed only when the armed call pushed, and only on that call's own output; any other input raises instead of reducing rows the exchange does not hold. Host state only (no GPU).""" import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_draft_lowm, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py index 50c75941815d..8bb087bf3cd4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -18,7 +18,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_draft_step as draft_step_kernels, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py index e7c9e5c19e6c..f8ffd8e47635 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The draft step's tile pieces and their summation order (kernels/m3_draft_step: the fixed warp order). +"""The draft step's tile pieces and their summation order (ops/m3_draft_step: the fixed warp order). A CTA's (tile, chunk) range is cut into 16 equal warp segments; the shared-memory atomics add a tile's pieces in arrival order, which is deterministic only while no tile takes more than 2 pieces. ``fixed_order`` sums them in warp @@ -14,7 +14,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( m3_draft_step as kds, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_step import ( # noqa: E501 @@ -43,7 +43,7 @@ @pytest.mark.parametrize("grid", [SMS, kds.GRID]) def test_real_grid_pieces(grid) -> None: - """One CTA per SM and the launch the target uses (kernels/m3_draft_step.GRID).""" + """One CTA per SM and the launch the target uses (ops/m3_draft_step.GRID).""" pieces = max_tile_pieces(DRAFT_SHARD // 16, grid, CHUNKS) assert (pieces > 2) == ORDER_DEPENDENT, (grid, pieces) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py index b865de3646c1..ef33aecf3a6a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -23,7 +23,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 @@ -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_107__tp4/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py index cf2b42bc86d9..d3e117c326e7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the forward-tail spec fold (kernels/m3_fold_tail and M3SpecStepFold.in_forward). +"""GPU test for the forward-tail spec fold (ops/m3_fold_tail and M3SpecStepFold.in_forward). The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers (step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing @@ -15,7 +15,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( modeling as spec_fold, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( m3_fold_tail, m3_spec_tail, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py index 9c7930e2a800..01aa7eb2af53 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (kernels/m3_lowm_splitk, -kernels/m3_draft_lowm): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (ops/m3_lowm_splitk, +ops/m3_draft_lowm): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's linears, eager and in a CUDA graph, with and without PDL.""" import pytest @@ -9,14 +9,14 @@ from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.m3_lowm_splitk import ( # noqa: E501 - SplitKTactic, - run_splitk_dense, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 ROWS, tactic_for, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, +) assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py index 51fee0cff1c5..4ed33c8a7251 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py @@ -174,7 +174,7 @@ def test_guard_refuses_fp64_temperatures() -> None: def test_the_verify_calls_this_targets_op_pair() -> None: """The verify calls this target's op pair (the sm_107 tp4 kernel), whose workspace sizes as the shared one.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_sampled_accept as kernels, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py index 7b0535e70812..8a70e7e9947a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py @@ -18,7 +18,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 triton_sparse_decode as tsd, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm import ( # noqa: E501 @@ -116,7 +116,7 @@ def test_the_dense_layers_take_the_fullctx_decode_at_the_measured_tp4_shapes() - def test_the_fullctx_plans_keep_the_tp4_geometry() -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py index 2c76156860e8..3046e9bc61bd 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py @@ -19,15 +19,12 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_index_score as index_score_kernel, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_index_score_t2sc as index_score_t2sc_kernel, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels.mx_gemv import ( # noqa: E501 - op as mx_gemv, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score import ( # noqa: E501 m3_index_score, ) @@ -36,6 +33,9 @@ SHARED_DOWN, SHARED_GATE_UP, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.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_107__tp4/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py index cd7d800bd1ca..0b5d60b19523 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py @@ -14,7 +14,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_mc_allreduce as mc, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py index 6fa19e3352fd..174be4bef5a1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test (SM 10.7) for kernels/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +"""GPU test (SM 10.7) for ops/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row first draft pass and the 1-row later passes of 1-2 requests. @@ -20,10 +20,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_drafter_attn as drafter, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_fullctx_decode as fullctx, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py index 03fc2b3576b3..b1759f370b7a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -16,10 +16,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_drafter_attn as drafter, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_fullctx_decode as fullctx, ) from tensorrt_llm._torch.attention.backends.fmha.utils import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py index 859e28a464d8..48d1bba922ac 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py @@ -112,7 +112,7 @@ def test_the_cells_cover_the_target_rows() -> None: def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) @@ -152,7 +152,7 @@ def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode step first hit (NaN rows from exp2(-inf - -inf)).""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl @@ -186,7 +186,7 @@ def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_fullctx_decode as kernel, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py index baf8ef7a2b73..b883a15acd38 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """The per-token valid-block count is written only before the step's first layer. -The long-row select (``kernels/m3_select_long_dsl.py``, ``SELECT_N_VALID_PREWAIT``) reads its row's entry of +The long-row select (``ops/m3_select_long_dsl.py``, ``SELECT_N_VALID_PREWAIT``) reads its row's entry of ``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the buffer happens before the step's first layer: @@ -131,7 +131,7 @@ def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: for func, line in _write_sites(ast.parse(src)): assert (rel, func) in _ALLOWED, ( f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " - "wait (kernels/m3_select_long_dsl.py, SELECT_N_VALID_PREWAIT), so a writer inside the " + "wait (ops/m3_select_long_dsl.py, SELECT_N_VALID_PREWAIT), so a writer inside the " "step must be ordered before the select's grid or the read moved after the wait" ) found.add((rel, func)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py index d5af7b1a4668..4022b2cfb133 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( m3_select_long_dsl as select_kernels, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long import ( # noqa: E501 @@ -140,7 +140,7 @@ def test_trigger_phase_selects_the_same(cell: int) -> None: @pytest.mark.parametrize("early_trigger", [False, True]) @pytest.mark.parametrize("cell", range(len(op.CELLS))) def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: - """kernels/m3_select_long_dsl.SELECT_N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the + """ops/m3_select_long_dsl.SELECT_N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the grid wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and tied rows).""" phase = 0 if early_trigger else 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py index 3883cd543732..36dbdd14d592 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py @@ -17,7 +17,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( modeling, # noqa: E501 ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_attn_lc8, m3_prequant, m3_report, @@ -594,7 +594,7 @@ def test_lc_fails_closed_on_ids_past_the_row(requests) -> None: """A selected id past the token's row (a stale table, e.g. a padded row keeping a 500K request's n_valid) is skipped, as the Triton producer skips it: the rows equal those with that slot at -1, bit for bit, and stay finite even though the block-table tail it would reach points at a page whose block scales are NaN (0xFF).""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_prequant, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py index cc8de6c62717..ac61687d6968 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py @@ -240,7 +240,7 @@ def test_run_to_run_deterministic(prewait, mxfp8) -> None: ids=["1x165k", "2req", "short-1row", "4req-incl-64"], ) def test_merge_once_is_bitwise_the_image_op(seq_lens, decode_query_len) -> None: - """kernels/m3_attn_lc8.MERGE_ONCE: the target's op (merge weights formed once per CTA) writes exactly the + """ops/m3_attn_lc8.MERGE_ONCE: the target's op (merge weights formed once per CTA) writes exactly the image op's rows and hand-off, at both softmax recurrences and both staging variants, including requests with fewer than 16 selected blocks (empty chunks, -inf LSE). With SPARSE_DECODE_LC8_PREFETCH_OPROJ it also pulls two ranges the size of the TP8 o_proj weight and scales (one not a multiple of 16 bytes) into L2: the same bits again. With @@ -318,7 +318,7 @@ def test_merge_once_is_bitwise_the_image_op(seq_lens, decode_query_len) -> None: def test_one_div_merge_is_deterministic_and_in_band( seq_lens, decode_query_len, prewait, linear ) -> None: - """kernels/m3_attn_lc8.ONE_DIV_MERGE (numerics rule 1b, gate a): the one-division merge writes the same rows on + """ops/m3_attn_lc8.ONE_DIV_MERGE (numerics rule 1b, gate a): the one-division merge writes the same rows on every run (5 eager runs, PDL on and off) and under CUDA-graph replay (3 replays == eager), MXFP8 hand-off included; the rows are within the float64 reference's band, and within summation-order distance of the per-chunk merge.""" case = build_step( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py index 769a903d51c2..e21326054b2e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test (SM 10.7) for kernels/m3_tokmap_guard: a one-request multi-token call of +"""GPU test (SM 10.7) for ops/m3_tokmap_guard: a one-request multi-token call of thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for @@ -16,7 +16,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_tokmap_guard as guard, ) from tensorrt_llm.bindings.internal import thop diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py index a9d64f3a87f7..c29fd0085966 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py @@ -7,15 +7,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__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_107__tp8.ops.attention.msa_kv_update import ( # noqa: E501 TOKENS_PER_BLOCK, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.msa_kv_update import ( # noqa: E501 msa_kv_update as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py index abf7fe23ea60..6d7def71212a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -358,7 +358,7 @@ def _capture_reference( updated: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor ) -> None: """The capture the boundary replaces: m3_spec_tail.capture_norm of the updated residual (no second addend).""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_spec_tail, ) @@ -674,7 +674,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_107__tp8.kernels.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_mc_allreduce as _ar # noqa: E501 from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm import ( # noqa: E501 m3_mc_allreduce_norm as _plain, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py index 50e69fc4b290..95374cf56848 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py @@ -276,8 +276,8 @@ def _run_one_rank() -> int: RANK, WORLD = rank, world assert WORLD == 8, f"the TP8 push is certified at 8 ranks, got {WORLD}" - import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_mc_allreduce as _ar # noqa: E501 - import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_prequant as _pq # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_prequant as _pq # noqa: E501 from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 sf_written as _sf_written, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py index d1793092cf8b..659b302d73ba 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py @@ -10,7 +10,7 @@ import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.mx_gemv import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv import ( # noqa: E501 op as gemv, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py index 1d1ffee82db5..1826394596eb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Bitwise oracle of the mx_gemv swigluoai prologue arithmetic (kernels/mx_gemv/_swiglu_oai_bf16x2) against the +"""Bitwise oracle of the mx_gemv swigluoai prologue arithmetic (ops/mx_gemv/_swiglu_oai_bf16x2) against the Triton silu_and_mul kernel the dense MLP runs today (modules/swiglu.py, alpha 1.702, limit 7, beta 1): every bf16 gate bit pattern against special and random up values, and every up bit pattern against special and random gates.""" @@ -10,7 +10,7 @@ import torch from cutlass.cute.runtime import from_dlpack -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.mx_gemv.m3_mx_gemv_kernel import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv.m3_mx_gemv_kernel import ( # noqa: E501 SWIGLU_OAI_ALPHA, SWIGLU_OAI_LIMIT, _swiglu_oai_bf16x2, @@ -121,7 +121,7 @@ def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E402 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.mx_gemv import ( # noqa: E402, E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv import ( # noqa: E402, E501 op as gemv, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py index f207d8896427..5a4dd5928876 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Rule 1b gate a for the TP8 decode GEMV order override (kernels/mx_gemv/op.STRIDED_K_TILES: strided k-tiles): +"""Rule 1b gate a for the TP8 decode GEMV order override (ops/mx_gemv/op.STRIDED_K_TILES: strided k-tiles): for every overridden (projection, rows) the output is the same on 3 eager calls and on a CUDA graph replay (deterministic), and its max abs / max rel error vs the fp64 product of the dequantized operands stays within the shipped order's band (the shared table's config on the same operands, over several seeds).""" @@ -8,9 +8,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.mx_gemv import ( # noqa: E501 - op, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 DENSE_DOWN, DENSE_GATE_UP, @@ -22,6 +19,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 mxfp8_dequant, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv import ( # noqa: E501 + op, +) pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py index 3299a8432362..7afbbaa79e26 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py @@ -12,7 +12,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_mc_allreduce, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py index 9a82b7505885..5a406341fa57 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py @@ -4,13 +4,13 @@ The draft-step and sampled-accept workspaces are created collectively (both ranks, a multicast buffer) on the first eager call. Creating one inside a capture would allocate and synchronize mid-capture, so a capture that finds it missing -reports a miss (``kernels/m3_report.capture_missed``) and gets None (the caller then takes its stock path) instead. +reports a miss (``ops/m3_report.capture_missed``) and gets None (the caller then takes its stock path) instead. """ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( m3_draft_step, m3_report, m3_sampled_accept, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py index 3bcee7c3345a..f07eecb1f0c2 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test of the fc fork's counter wait (modeling's DI_FC_FLAG, kernels/m3_fc_wait): the one-thread wait kernel +"""GPU test of the fc fork's counter wait (modeling's DI_FC_FLAG, ops/m3_fc_wait): the one-thread wait kernel queued on the main stream releases its successor only after the fork's bump on another stream (so the successor reads what the fork wrote, with no stream join before it), eagerly and on CUDA-graph replays; it counts an overrun instead of hanging when no bump comes.""" @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_fc_wait, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py index e1b7c9c625bb..f1f5f6160e32 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py @@ -6,7 +6,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_spec_tail, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_input import ( # noqa: E501 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py index 58e3c0a89b6d..877e52628100 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -18,7 +18,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_draft_step as draft_step_kernels, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_target_head import ( # noqa: E501 @@ -216,7 +216,7 @@ def _args( def _fixed(args) -> bool: - """fixed_order for these op arguments as dispatch picks it (kernels/m3_draft_step._fixed_order): the warp-order + """fixed_order for these op arguments as dispatch picks it (ops/m3_draft_step._fixed_order): the warp-order sum where the launch grid lets a tile take more than 2 warp pieces (D-FIX; the TP8 shard on 212 SMs does).""" x, wf, num_ctas, exc_off = args[0], args[3], args[20], args[23] tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 @@ -698,7 +698,7 @@ def test_guard_refuses_too_few_ctas() -> None: op(*args, fixed_order=_fixed(args)) -# ------------------------------------------------- one pass as two launches (uGPU, kernels/m3_draft_step.UGPU) +# ------------------------------------------------- one pass as two launches (uGPU, ops/m3_draft_step.UGPU) @pytest.mark.parametrize( @@ -710,7 +710,7 @@ def test_two_launches_per_domain_are_the_one_launch_pass(cell) -> None: """A pass split into two launches (tiles [0, split) and [split, n), one per locality domain stream, the exchange and bookkeeping by the last of both launches' CTAs) writes exactly what the one-launch pass writes: tokens, winning keys, draft columns, next positions, KV and sequence lengths, and re-arms the counters and keys.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_locality import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_locality import ( # noqa: E501 LocalityMap, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py index 601b65e50ed2..ec29cf2a88e7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""D-FIX: the draft step's tile pieces and their summation order (kernels/m3_draft_step.FIXED_ORDER). +"""D-FIX: the draft step's tile pieces and their summation order (ops/m3_draft_step.FIXED_ORDER). A CTA's (tile, chunk) range is cut into 16 equal warp segments; the shared-memory atomics add a tile's pieces in arrival order, which is deterministic only while no tile takes more than 2 pieces. ``fixed_order`` sums them in warp @@ -14,7 +14,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( m3_draft_step as kds, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_step import ( # noqa: E501 @@ -43,7 +43,7 @@ @pytest.mark.parametrize("grid", [SMS, kds.GRID]) def test_real_grid_pieces(grid) -> None: - """One CTA per SM and the DS-GRID launch (kernels/m3_draft_step.GRID).""" + """One CTA per SM and the DS-GRID launch (ops/m3_draft_step.GRID).""" pieces = max_tile_pieces(DRAFT_SHARD // 16, grid, CHUNKS) assert (pieces > 2) == ORDER_DEPENDENT, (grid, pieces) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py index e6d2fb864da9..5e6ae369f8fa 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the forward-tail spec fold (kernels/m3_fold_tail and M3SpecStepFold.in_forward). +"""GPU test for the forward-tail spec fold (ops/m3_fold_tail and M3SpecStepFold.in_forward). The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers (step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing @@ -15,7 +15,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( m3_fold_tail, m3_spec_tail, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py index b66a1e35e4a4..821c6eeb08ae 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (kernels/m3_lowm_splitk, +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (ops/m3_lowm_splitk, modeling.DRAFT_LOWM_OWN_SPLITK): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's linears, eager and in a CUDA graph, with and without PDL.""" @@ -9,14 +9,14 @@ from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels.m3_lowm_splitk import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + TACTICS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_lowm_splitk import ( # noqa: E501 SplitKTactic, run_splitk_dense, run_splitk_dense_swiglu, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 - TACTICS, -) assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py index f4739cf8f5fc..6066e23c92a3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py @@ -131,7 +131,7 @@ def _walk_outputs(args, one_round_k, ilp, calls=3, offset_step=1000): @pytest.mark.parametrize("one_round_k", [0, 82]) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_the_ilp_walk_changes_nothing(cell, one_round_k) -> None: - """kernels/m3_sampled_accept.ILP_WALK (SA-PW) only computes the owner's token weights 8 at a time; the walk adds + """ops/m3_sampled_accept.ILP_WALK (SA-PW) only computes the owner's token weights 8 at a time; the walk adds them in the same order with the same early exit, so the draws and acceptance outputs are the one-token walk's, bit for bit, call after call, with SA-1R off and on.""" one = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py index 914e0d38bf05..3925faf6082a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""SA-1R (kernels/m3_sampled_accept.ONE_ROUND_K) on trtllm::minimax_m3_sm107tp8_sampled_accept, world 1, the TP8 shard. +"""SA-1R (ops/m3_sampled_accept.ONE_ROUND_K) on trtllm::minimax_m3_sm107tp8_sampled_accept, world 1, the TP8 shard. The one-round variants (one_round_k 9: exact top-9 list; 82: pooled top-8 + bound) must return the two-round kernel's (one_round_k 0) sampled tokens, accepted tokens, accepted counts, gather ids and extra-attempt count BIT FOR diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py index fd25081cfe12..fe4a08348fb1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py @@ -16,10 +16,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_draft_step as draft_step_kernel, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 triton_sparse_decode as tsd, ) @@ -226,7 +226,7 @@ def test_the_sparse_decode_writes_the_mxfp8_handoff() -> None: def test_sparse_decode_lc8_covers_the_decode_steps_with_the_handoff() -> None: """Every decode step (at most DECODE_MAX_TOKENS rows) writes the MXFP8 o_proj hand-off; the one-launch decode stages its pre-wait reads up to m3_attn_lc8.PREWAIT_MAX_ROWS rows and merges once.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_attn_lc8, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py index 9414fae5c97f..8f28a67e722f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py @@ -14,7 +14,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.kernels import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_mc_allreduce as mc, ) 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 bfb70013cd4e..9fc45ae71a11 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -715,6 +715,39 @@ def test_minimax_m3_route_declines_another_rope_base_or_sparse_layout(monkeypatc ) +@pytest.mark.parametrize("mapping_kwargs", [_M3_TP2, _M3_TP4, _M3_TP8], ids=["tp2", "tp4", "tp8"]) +@pytest.mark.parametrize( + "layout", + [ + _OTHER_LAYOUT, + dict( + _M3_LAYOUT, + sparse_attention_config=dict( + _M3_LAYOUT["sparse_attention_config"], use_sparse_attention=False + ), + ), + ], + ids=["other_revision", "sparse_disabled"], +) +def test_minimax_m3_rubin_route_declines_another_layout(monkeypatch, mapping_kwargs, layout): + """The sm107 targets build the same RoPE table and sparse layout as the sm103 ones, so they decline the same + checkpoints.""" + from types import SimpleNamespace + + from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + with pytest.raises(NoTarget, match="sparse attention layers other than 3-59|rope_theta"): + routing.route( + SimpleNamespace(text_config=_m3_text_config(**layout)), + Mapping(**mapping_kwargs), + routing._SM_RUBIN, + None, + QuantAlgo.NVFP4, + ) + + 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 From 076590dbe45f7f14b0b3533b3ac9c7e66e318435 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 21 Oct 2026 18:06:25 -0700 Subject: [PATCH 15/29] [None][fix] Run the MiniMax-M3 E4M3 cast probe on the CPU under any default device The import-time probe built its tensors without a device. When the module is first imported while a model is built under the meta device, the probe tensors are meta and torch.equal raises. Pin them to the CPU, as the comment already states. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../backends/sparse/minimax_m3/kernels/paged_cache.py | 7 +++++-- .../attention/test_modeling_v2_m3_fp8_saturate.py | 8 ++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py index 6504f0fbd2d4..2f40c6b090f2 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py @@ -19,8 +19,11 @@ def _cast_saturates() -> bool: """Whether this torch's E4M3 cast maps every out-of-range value, infinities included, to +-448.""" - probe = torch.tensor([449.0, 463.0, 464.0, 1e4, float("inf"), -464.0, -float("inf")]) - want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0]) + # An explicit device: the first import can run under a default-device context (meta while a model is built). + probe = torch.tensor( + [449.0, 463.0, 464.0, 1e4, float("inf"), -464.0, -float("inf")], device="cpu" + ) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0], device="cpu") return torch.equal(probe.to(torch.float8_e4m3fn).float(), want) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py index f58dc5f3f0ce..adfd40f1d941 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py @@ -100,3 +100,11 @@ def test_fp8_e4m3_saturate_takes_the_cast_only_path_here(monkeypatch) -> None: assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 1 monkeypatch.setattr(paged_cache, "_CAST_SATURATES", False) assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 2 + + +@pytest.mark.parametrize("device", ["meta", "cuda"]) +def test_cast_probe_ignores_the_default_device(device) -> None: + """The import-time probe runs on the CPU even under a default-device context, as when the module is first + imported while a model is built on the meta device.""" + with torch.device(device): + assert paged_cache._cast_saturates() == paged_cache._CAST_SATURATES 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 16/29] [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 3dac8aedcdcbad0b0eec9916863eaffd16f8f7a2 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:34:03 -0700 Subject: [PATCH 17/29] [None][test] Move the MXFP8 CuTe pinned-tactics test next to the dense MXFP8 GEMM tests Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../{moe => }/kernels/parallel/test_mxfp8_cute_pinned_tactics.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/unittest/_torch/{moe => }/kernels/parallel/test_mxfp8_cute_pinned_tactics.py (100%) diff --git a/tests/unittest/_torch/moe/kernels/parallel/test_mxfp8_cute_pinned_tactics.py b/tests/unittest/_torch/kernels/parallel/test_mxfp8_cute_pinned_tactics.py similarity index 100% rename from tests/unittest/_torch/moe/kernels/parallel/test_mxfp8_cute_pinned_tactics.py rename to tests/unittest/_torch/kernels/parallel/test_mxfp8_cute_pinned_tactics.py 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 18/29] [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 19/29] [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] From 8315bd022069a8dd50b983bafc9988d45d92a16b Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 27 Oct 2026 15:21:27 -0700 Subject: [PATCH 20/29] [None][fix] Scope context warmup to Rubin and pin stock M3 coverage Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../kernels/trtllmGenKernels/fmha/fmhaKernels.h | 6 ++++++ tests/integration/defs/accuracy/test_llm_api_pytorch.py | 2 ++ 2 files changed, 8 insertions(+) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 02e50f852af6..ef5aa4a6ad5b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -530,6 +530,12 @@ class TllmGenFmhaKernel } if (isContextKernel(runnerParams.mKernelType)) { + // Only SM107 needs the dense paged-context NVRTC warmup grid. + // Other architectures retain their precompiled context path. + if (mSM != kSM_107) + { + return; + } cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; TLLM_CUDA_CHECK(cudaStreamIsCapturing(runnerParams.stream, &captureStatus)); TLLM_CHECK_WITH_INFO(captureStatus == cudaStreamCaptureStatusNone, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 9211e512e986..69e26467d31b 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7589,6 +7589,7 @@ def _run_nvfp4(self, torch_compile_config=TorchCompileConfig() if piecewise else None, max_seq_len=4096, + modeling_v2="off", trust_remote_code=True) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION if kv_dtype == "nvfp4": @@ -7712,6 +7713,7 @@ def _run_nvfp4_eagle3(self, # Keep the whole acceptance probe: the default buffer holds # only the latest 1000 iterations. max_stats_len=5000, + modeling_v2="off", trust_remote_code=True, **piecewise_kwargs) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION From d1d824814b189b2d7eacd97566872a9a150ea698 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 27 Oct 2026 15:28:11 -0700 Subject: [PATCH 21/29] [None][test] Complete M3 block-table comparison fixture Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/attention/sparse/msa/test_msa_runtime_policy.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 d0e6d45179aa..ec98dc90449b 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 @@ -412,7 +412,10 @@ def _page_index_views(seed: int): def _bare_manager() -> MiniMaxM3KVCacheManagerV2: """A manager with only the state get_block_ids_per_seq reads (no pools, no V2 base).""" - return object.__new__(MiniMaxM3KVCacheManagerV2) + manager = object.__new__(MiniMaxM3KVCacheManagerV2) + manager.layer_specs = None + manager.kv_factor = 2 + return manager def test_block_ids_from_views_is_off_by_default(): From 5f4c412ddacad341a8e3ff60cdc5a1d0c1cc4d2f Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:20:16 -0700 Subject: [PATCH 22/29] [None][fix] Use decoder callback for M3 graph headroom Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py | 2 +- .../_torch/executor/engine/test_decoder_runtime_hooks.py | 2 +- 2 files changed, 2 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 4abeb8c1a193..e50173a1f940 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -541,7 +541,7 @@ def _short_long_seq_len_threshold(self) -> int | None: f"attention config, got {type(sparse_config).__name__}" ) topk_span = sparse_config.sparse_topk_blocks * sparse_config.sparse_block_size - return topk_span - self.get_runtime_tokens_per_gen_step(self._config.max_draft_len) + return topk_span - self._get_runtime_tokens_per_gen_step(self._config.max_draft_len) return None def _initialize_breakable_cuda_graph_runner(self) -> BreakableCUDAGraphRunner | None: diff --git a/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py index 9aeb6df53c14..f6ccb731c0cf 100644 --- a/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py +++ b/tests/unittest/_torch/executor/engine/test_decoder_runtime_hooks.py @@ -41,7 +41,7 @@ def _runner( runner._sparse_attention_config = sparse_attention_config runner.runtime_hooks = hooks runner._config = SimpleNamespace(max_draft_len=max_draft_len) - runner.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 + runner._get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 return runner From 890400f44e891ba23153ced724538a35856742cd Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:04:57 -0700 Subject: [PATCH 23/29] [None][fix] Remove duplicate M3 CPU test registration Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_cpu.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 449e6ca4d970..0d8b73ae0a26 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -13,7 +13,6 @@ l0_cpu: orchestrator: mpi tests: # MiniMax-M3 Rubin target unit tests. - - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py # MiniMax-M3 target unit tests. - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py From 7362f43bace852e2a25876b4b3c17b442c426059 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:35:06 -0700 Subject: [PATCH 24/29] [None][fix] Isolate M3 fake registrations from stock imports Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../modeling_v2/models/minimax_m3/__init__.py | 2 -- .../modeling_v2/models/minimax_m3/fake_ops.py | 14 ++++---------- .../minimax_m3_nvfp4__sm_103__tp4/__init__.py | 2 ++ .../ops/m3_attn_mega.py | 6 ++++-- .../ops/m3_attn_mega_c1.py | 6 ++++-- .../ops/m3_fused_attn.py | 6 ++++-- .../minimax_m3_nvfp4__sm_103__tp8/__init__.py | 2 ++ tensorrt_llm/_torch/pyexecutor/step_inputs.py | 9 ++++++--- 8 files changed, 26 insertions(+), 21 deletions(-) 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 5c445343d34d..a3dfe4a6ce39 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,5 +1,3 @@ # 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 index c47398888520..a031ea7f9e70 100644 --- 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 @@ -1,17 +1,11 @@ # 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 +"""Fake implementations of the MiniMax-M3 decode ops for compile and fake-tensor tracing. +Target packages register these implementations when imported after routing. +""" -@torch.library.register_fake("trtllm::step_inputs_scatter") -def _(stage): - # Writes the staged rows into their destinations; returns nothing. - return None +import torch @torch.library.register_fake("trtllm::minimax_m3_mxfp8_low_m_gemm") 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 index fdb324a42eda..33b2434f5049 100644 --- 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 @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_103 / tp4.""" + +from .. import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 index 964a138afd67..7c5550387f7e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 @@ -67,11 +67,13 @@ 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"): + try: + torch.ops.trtllm.minimax_m3_attn_mega + except AttributeError as exc: raise RuntimeError( "torch.ops.trtllm.minimax_m3_attn_mega is not registered: this TensorRT-LLM build lacks " "cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu" - ) + ) from exc _module = torch.ops.trtllm return _module diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 index 82c3b1cc3593..2315fbe12afd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 @@ -78,11 +78,13 @@ 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"): + try: + torch.ops.trtllm.minimax_m3_attn_c1 + except AttributeError as exc: raise RuntimeError( "torch.ops.trtllm.minimax_m3_attn_c1 is not registered: this TensorRT-LLM build lacks " "cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu" - ) + ) from exc _module = torch.ops.trtllm return _module diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 index fefd30acf333..e4bd922a1587 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 @@ -75,11 +75,13 @@ 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"): + try: + torch.ops.trtllm.minimax_m3_fused_decode_attn + except AttributeError as exc: 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" - ) + ) from exc _module = torch.ops.trtllm return _module 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 index 2e340ff7284d..f7e603e6af6f 100644 --- 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 @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_103 / tp8.""" + +from .. import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/pyexecutor/step_inputs.py b/tensorrt_llm/_torch/pyexecutor/step_inputs.py index ff0c02a95611..c538896c04da 100644 --- a/tensorrt_llm/_torch/pyexecutor/step_inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -29,12 +29,15 @@ 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 +@torch.library.register_fake("trtllm::step_inputs_scatter") +def _step_inputs_scatter_fake(stage: torch.Tensor) -> None: + # Writes the staged rows into their destinations; returns nothing. + return None + + class StepInputPack: """One engine's pack of deferred step-input copies (see the module docstring).""" From 54da95bfd9b76463f06557c571ae23057a8f491a Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:36:58 -0700 Subject: [PATCH 25/29] [None][fix] Isolate M3 fake registrations from stock imports Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../modeling_v2/models/minimax_m3/__init__.py | 2 -- .../modeling_v2/models/minimax_m3/fake_ops.py | 14 ++++---------- .../minimax_m3_nvfp4__sm_103__tp4/__init__.py | 2 ++ .../ops/m3_attn_mega.py | 6 ++++-- .../ops/m3_attn_mega_c1.py | 6 ++++-- .../ops/m3_fused_attn.py | 6 ++++-- .../minimax_m3_nvfp4__sm_103__tp8/__init__.py | 2 ++ .../minimax_m3_nvfp4__sm_107__tp2/__init__.py | 2 ++ .../minimax_m3_nvfp4__sm_107__tp4/__init__.py | 2 ++ .../minimax_m3_nvfp4__sm_107__tp8/__init__.py | 2 ++ tensorrt_llm/_torch/pyexecutor/step_inputs.py | 9 ++++++--- 11 files changed, 32 insertions(+), 21 deletions(-) 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 5c445343d34d..a3dfe4a6ce39 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,5 +1,3 @@ # 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 index 7576ce1fe445..8e589f83d1c0 100644 --- 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 @@ -1,17 +1,11 @@ # 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 +"""Fake implementations of the MiniMax-M3 decode ops for compile and fake-tensor tracing. +Target packages register these implementations when imported after routing. +""" -@torch.library.register_fake("trtllm::step_inputs_scatter") -def _(stage): - # Writes the staged rows into their destinations; returns nothing. - return None +import torch @torch.library.register_fake("trtllm::minimax_m3_mxfp8_low_m_gemm") 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 index fdb324a42eda..33b2434f5049 100644 --- 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 @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_103 / tp4.""" + +from .. import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 index 964a138afd67..7c5550387f7e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 @@ -67,11 +67,13 @@ 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"): + try: + torch.ops.trtllm.minimax_m3_attn_mega + except AttributeError as exc: raise RuntimeError( "torch.ops.trtllm.minimax_m3_attn_mega is not registered: this TensorRT-LLM build lacks " "cpp/tensorrt_llm/kernels/minimaxM3/attnMega.cu" - ) + ) from exc _module = torch.ops.trtllm return _module diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 index 82c3b1cc3593..2315fbe12afd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 @@ -78,11 +78,13 @@ 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"): + try: + torch.ops.trtllm.minimax_m3_attn_c1 + except AttributeError as exc: raise RuntimeError( "torch.ops.trtllm.minimax_m3_attn_c1 is not registered: this TensorRT-LLM build lacks " "cpp/tensorrt_llm/kernels/minimaxM3/attnC1.cu" - ) + ) from exc _module = torch.ops.trtllm return _module diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 index fefd30acf333..e4bd922a1587 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/ops/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 @@ -75,11 +75,13 @@ 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"): + try: + torch.ops.trtllm.minimax_m3_fused_decode_attn + except AttributeError as exc: 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" - ) + ) from exc _module = torch.ops.trtllm return _module 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 index 2e340ff7284d..f7e603e6af6f 100644 --- 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 @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_103 / tp8.""" + +from .. import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py index 25c6fb9961de..d1979a7761d7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_107 / tp2.""" + +from .. import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py index 36a53a834458..cca058a4265e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_107 / tp4.""" + +from .. import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py index adf71d21a10a..a2bc9cac05c9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_107 / tp8.""" + +from .. import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/pyexecutor/step_inputs.py b/tensorrt_llm/_torch/pyexecutor/step_inputs.py index 29b3eb318afd..448064f4968c 100644 --- a/tensorrt_llm/_torch/pyexecutor/step_inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -61,9 +61,6 @@ import triton import triton.language as tl -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 @@ -148,6 +145,12 @@ def _parallel_ok(h2d: list[tuple[int, int, int]], d2d: list[tuple[int, int, int] return all(_disjoint(dsts + [(src, n)]) for _, src, n in d2d) +@torch.library.register_fake("trtllm::step_inputs_scatter") +def _step_inputs_scatter_fake(stage: torch.Tensor) -> None: + # Writes the staged rows into their destinations; returns nothing. + return None + + class StepInputPack: """One engine's pack of deferred step-input copies (see the module docstring).""" From 0b1a818bf3c0de92237f22573d80a4b1f87aa9cd Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:41:19 -0700 Subject: [PATCH 26/29] [None][fix] Require compiled Rubin target operators Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../minimax_m3_nvfp4__sm_107__tp2/modeling.py | 17 +++--------- .../ops/m3_fullctx_decode.py | 5 ++-- .../minimax_m3_nvfp4__sm_107__tp4/modeling.py | 7 +---- .../ops/m3_fullctx_decode.py | 9 ++++--- .../minimax_m3_nvfp4__sm_107__tp8/modeling.py | 17 +++--------- .../ops/m3_fullctx_decode.py | 26 ++++++++----------- .../ops/m3_mc_allreduce.py | 18 ++++--------- 7 files changed, 31 insertions(+), 68 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py index 2d02fac77dcf..4418d2908b33 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py @@ -2154,12 +2154,7 @@ def _qk_norm_rope(core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool) def _dense_fp8_op_applies(qkv: torch.Tensor, position_ids) -> bool: """Whether a dense layer's FP8 q/k/v come from this target's op (``minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8``, else the BF16 fused_qk_norm_rope output cast to FP8).""" - return ( - DENSE_QKV_FP8_VIA_BF16 - and hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8") - and position_ids is not None - and qkv.dtype == torch.bfloat16 - ) + return DENSE_QKV_FP8_VIA_BF16 and position_ids is not None and qkv.dtype == torch.bfloat16 def _dense_qk_norm_rope_fp8( @@ -2167,7 +2162,7 @@ def _dense_qk_norm_rope_fp8( ) -> torch.Tensor: """A dense layer's FP8 q/k/v: QK norm + RoPE with Q/K rounded to BF16 and then to E4M3 and V cast, the bits of the BF16 fused_qk_norm_rope output cast to FP8 (K is stored in the FP8 cache), in one launch of this target's op - (``DENSE_QKV_FP8_VIA_BF16``); without the op, the BF16 kernel and the cast. + (``DENSE_QKV_FP8_VIA_BF16``); outside its supported inputs, the BF16 kernel and the cast. With ``kv_dest`` = (the layer's HND pool view, the step's int32 cache slots, live tokens, a [tokens, q heads, 128] FP8 buffer) the op also stores the live tokens' K and V rows into the pool (the bytes ``write_msa_phase_kv`` @@ -4018,7 +4013,6 @@ def _dense_attention(self, i: int, position_ids, hidden, step, push): and fp8 and DENSE_KV_INSERT_FUSED and _dense_fp8_op_applies(qkv, position_ids) - and hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") ): kv_dest = ( step["kv_cache_manager"].get_buffers(i, kv_layout="HND"), @@ -6358,12 +6352,7 @@ def cache_derived_state(self) -> None: FULLCTX_DECODE_MAX_KV_LEN, FULLCTX_DECODE_OWN_REDUCER and m3_fullctx_decode_kernel.has_own_reducer(), get_env_enable_pdl(), - fused_rope_mode=( - DRAFTER_ROPE_MODE - if DRAFTER_ROPE_FUSED - and hasattr(torch.ops.trtllm, "minimax_m3_sm107tp2_drafter_rope_fp8_kv") - else None - ), + fused_rope_mode=(DRAFTER_ROPE_MODE if DRAFTER_ROPE_FUSED else None), on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), on_taken=lambda rows: taken(path, rows), on_uncovered=lambda reason: uncovered(path, reason), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py index c220741e3874..b6705bdb46b7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py @@ -258,8 +258,9 @@ def _check_layout(cfg, qlen: int, splits: int) -> None: def has_own_reducer() -> bool: - """Whether this build carries the BF16 split-KV merge op.""" - return hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_fullctx_reduce") + """Require the BF16 split-KV merge op supplied by this target build.""" + torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce + return True def plan_for( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py index fbc60ddc1d76..550316d84740 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py @@ -5830,12 +5830,7 @@ def cache_derived_state(self) -> None: FULLCTX_DECODE_MAX_KV_LEN, FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), get_env_enable_pdl(), - fused_rope_mode=( - DRAFTER_ROPE_MODE - if DRAFTER_ROPE_FUSED - and hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_drafter_rope_fp8_kv") - else None - ), + fused_rope_mode=(DRAFTER_ROPE_MODE if DRAFTER_ROPE_FUSED else None), on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), on_taken=lambda rows: taken(path, rows), on_uncovered=lambda reason: uncovered(path, reason), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py index ef6b9ee0d1ad..3dc0fc4a3888 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py @@ -270,13 +270,14 @@ def _check_layout(cfg, qlen: int, splits: int) -> None: def emits_mxfp8() -> bool: """Whether the loaded merge op can also write the o_proj MXFP8 hand-off (its schema takes q8 / sf).""" - op = getattr(torch.ops.trtllm, "minimax_m3_sm107tp4_fullctx_reduce", None) - return op is not None and "q8" in str(op.default._schema) + op = torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce + return "q8" in str(op.default._schema) def has_own_reducer() -> bool: - """Whether this build carries the BF16 split-KV merge op.""" - return hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_fullctx_reduce") + """Require the BF16 split-KV merge op supplied by this target build.""" + torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce + return True def plan_for( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py index 490687087aee..69ea6f145b80 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py @@ -2018,12 +2018,7 @@ def _qk_norm_rope(core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool) def _dense_fp8_op_applies(qkv: torch.Tensor, position_ids) -> bool: """Whether a dense layer's FP8 q/k/v come from this target's op (``minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8``, else the BF16 fused_qk_norm_rope output cast to FP8).""" - return ( - DENSE_QKV_FP8_VIA_BF16 - and hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8") - and position_ids is not None - and qkv.dtype == torch.bfloat16 - ) + return DENSE_QKV_FP8_VIA_BF16 and position_ids is not None and qkv.dtype == torch.bfloat16 def _dense_qk_norm_rope_fp8( @@ -2031,7 +2026,7 @@ def _dense_qk_norm_rope_fp8( ) -> torch.Tensor: """A dense layer's FP8 q/k/v: QK norm + RoPE with Q/K rounded to BF16 and then to E4M3 and V cast, the bits of the BF16 fused_qk_norm_rope output cast to FP8 (K is stored in the FP8 cache), in one launch of this target's op - (``DENSE_QKV_FP8_VIA_BF16``); without the op, the BF16 kernel and the cast. + (``DENSE_QKV_FP8_VIA_BF16``); outside its supported inputs, the BF16 kernel and the cast. With ``kv_dest`` = (the layer's HND pool view, the step's int32 cache slots, live tokens, a [tokens, q heads, 128] FP8 buffer) the op also stores the live tokens' K and V rows into the pool (the bytes ``write_msa_phase_kv`` @@ -3720,7 +3715,6 @@ def _dense_attention(self, i: int, position_ids, hidden, step, push): and fp8 and DENSE_KV_INSERT_FUSED and _dense_fp8_op_applies(qkv, position_ids) - and hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") ): kv_dest = ( step["kv_cache_manager"].get_buffers(i, kv_layout="HND"), @@ -6045,12 +6039,7 @@ def cache_derived_state(self) -> None: FULLCTX_DECODE_MAX_KV_LEN, FULLCTX_DECODE_OWN_REDUCER and m3_fullctx_decode_kernel.has_own_reducer(), get_env_enable_pdl(), - fused_rope_mode=( - DRAFTER_ROPE_MODE - if DRAFTER_ROPE_FUSED - and hasattr(torch.ops.trtllm, "minimax_m3_sm107tp8_drafter_rope_fp8_kv") - else None - ), + fused_rope_mode=(DRAFTER_ROPE_MODE if DRAFTER_ROPE_FUSED else None), on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), on_taken=lambda rows: taken(path, rows), on_uncovered=lambda reason: uncovered(path, reason), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py index 71666eb05872..de9f6d6666f3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py @@ -268,34 +268,30 @@ def _check_layout(cfg, qlen: int, splits: int) -> None: raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) -# The split-KV merge runs the TP8 copy of the BF16 reducer when the op library carries it (the same bits). +# Select the TP8 copy of the BF16 reducer; both copies use the same per-row arithmetic. TP8_REDUCER = True def _reduce_op(): - """The BF16 split-KV merge: the TP8 copy (one row per CTA, its dependent released right after its wait) when the - TP8 op library carries it, else the image's TP2 op; the same per-row arithmetic, so the same bits.""" - ops = torch.ops.trtllm - if TP8_REDUCER and hasattr(ops, "minimax_m3_sm107tp8_fullctx_reduce"): - return ops.minimax_m3_sm107tp8_fullctx_reduce - return ops.minimax_m3_sm107_tp2_fullctx_reduce + """The BF16 split-KV merge selected by TP8_REDUCER.""" + if TP8_REDUCER: + return torch.ops.trtllm.minimax_m3_sm107tp8_fullctx_reduce + return torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce def _tp8_reduce_op(): - return getattr(torch.ops.trtllm, "minimax_m3_sm107tp8_fullctx_reduce", None) + return torch.ops.trtllm.minimax_m3_sm107tp8_fullctx_reduce def emits_mxfp8() -> bool: - """Whether the merge run uses can also write the o_proj MXFP8 hand-off (the TP8 copy of the merge).""" - return _tp8_reduce_op() is not None and _reduce_op() is _tp8_reduce_op() + """Whether the selected merge writes the TP8 o_proj MXFP8 hand-off.""" + return _reduce_op() is _tp8_reduce_op() def has_own_reducer() -> bool: - """Whether this build carries a BF16 split-KV merge op (the TP8 copy or the image's TP2 op).""" - ops = torch.ops.trtllm - return hasattr(ops, "minimax_m3_sm107tp8_fullctx_reduce") or hasattr( - ops, "minimax_m3_sm107_tp2_fullctx_reduce" - ) + """Require the selected BF16 split-KV merge op supplied by this target build.""" + _reduce_op() + return True def plan_for( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py index bb87d8ead5bd..eb025bc1694a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py @@ -319,7 +319,8 @@ def m3_mc_finalize_allreduce_norm( @functools.lru_cache(maxsize=None) def _own_op_loaded() -> bool: - return getattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_mc_allreduce_norm", None) is not None + torch.ops.trtllm.minimax_m3_sm107_tp8_mc_allreduce_norm + return True # The boundary all-reduces run this target's op (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm) rather than the @@ -329,17 +330,10 @@ def _own_op_loaded() -> bool: def own_op() -> bool: - """OWN_OP when this target's op is loaded (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm), with - its engagement line; else the image's op. The same library is loaded on every rank, so the ranks agree.""" + """Select this target's all-reduce when OWN_OP is enabled, requiring its compiled op.""" if not OWN_OP: return False - if not _own_op_loaded(): - m3_report.missed( - OWN_OP_PATH, - None, - "trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm is not loaded (the image's op runs)", - ) - return False + _own_op_loaded() m3_report.taken(OWN_OP_PATH, None, "the boundary all-reduces run this target's op") return True @@ -353,9 +347,7 @@ def capture_supported() -> bool: if own_op() else "minimax_m3_sm107_tp2_mc_allreduce_norm" ) - op = getattr(torch.ops.trtllm, name, None) - if op is None: - return False + op = getattr(torch.ops.trtllm, name) return any(a.name == "cap_out" for a in op.default._schema.arguments) From 2b23cdc49c1141e50766c80a7b12fd250cf7bdc0 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:35:01 -0700 Subject: [PATCH 27/29] [None][fix] Keep shared M3 fake ops in the kernel package Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../modeling_v2/models/minimax_m3/fake_ops.py | 176 +---------------- .../models/minimax_m3/kernels/__init__.py | 3 + .../models/minimax_m3/kernels/fake_ops.py | 177 ++++++++++++++++++ .../minimax_m3_nvfp4__sm_103__tp4/__init__.py | 2 +- .../minimax_m3_nvfp4__sm_103__tp8/__init__.py | 2 +- .../minimax_m3_nvfp4__sm_107__tp2/__init__.py | 2 +- .../minimax_m3_nvfp4__sm_107__tp4/__init__.py | 2 +- .../minimax_m3_nvfp4__sm_107__tp8/__init__.py | 2 +- 8 files changed, 187 insertions(+), 179 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py 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 index 8e589f83d1c0..76ea86dfc52e 100644 --- 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 @@ -1,177 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Fake implementations of the MiniMax-M3 decode ops for compile and fake-tensor tracing. +"""Register MiniMax-M3 fake ops through their shared kernel module.""" -Target packages register these implementations when imported after routing. -""" - -import torch - - -@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 - - -# The MiniMax-M3 sm_107 TP2 / TP4 / TP8 decode ops: one producer and one all-reduce + norm op per TP width (Q and -# index-Q heads per rank from the kernels' producerPw{2,4,8}.h), and the dense QK-norm + RoPE ops. -def _register_minimax_m3_sm107_fakes(tp: int, heads_q: int, heads_index: int) -> None: - @torch.library.register_fake(f"trtllm::minimax_m3_sm107_tp{tp}_producer_pw") - def _( - packed, - kv_data, - kv_scale, - index_k, - out_cache_loc, - kv_quant_scale, - eps, - q_weight, - k_weight, - index_q_weight, - index_k_weight, - rotary_cos_sin, - position_ids, - pdl, - early, - prewait_step_inputs=False, - ): - # [Q, index Q] in FP8; the K / V / index-K rows go into the caches in place. - num_tokens = packed.shape[0] - return [ - packed.new_empty((num_tokens, heads_q, 128), dtype=torch.float8_e4m3fn), - packed.new_empty((num_tokens, heads_index, 128), dtype=torch.float8_e4m3fn), - ] - - @torch.library.register_fake(f"trtllm::minimax_m3_sm107_tp{tp}_mc_allreduce_norm") - def _( - x, - x2, - residual, - norm_w, - eps, - ws_uc, - ws_mc_ptr, - flags, - rank, - world, - cluster, - pdl, - perm, - scales, - maxt, - quant, - prepushed, - opts, - cap_out=None, - cap_w=None, - cap_eps=0.0, - ): - # [normed, updated residual] (+ [E4M3 rows, swizzled scale bytes] with quant); cap_out is written in place. - outputs = [torch.empty_like(residual), torch.empty_like(residual)] - if quant: - m, h = residual.shape - outputs += [ - residual.new_empty((m, h), dtype=torch.float8_e4m3fn), - residual.new_empty(((m + 127) // 128 * 128 * (h // 32),), dtype=torch.uint8), - ] - return outputs - - -for _tp, _heads_q, _heads_index in ((2, 32, 2), (4, 16, 1), (8, 8, 1)): - _register_minimax_m3_sm107_fakes(_tp, _heads_q, _heads_index) - - -def _minimax_m3_sm107_dense_qk_norm_rope_fp8_fake(qkv, *args, **kwargs): - # A new FP8 tensor shaped like qkv; the KV variant also stores K / V (and the compact Q) in place. - return torch.empty_like(qkv, dtype=torch.float8_e4m3fn) - - -for _name in ( - "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8", - "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv", - "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8", - "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv", -): - torch.library.register_fake(f"trtllm::{_name}")(_minimax_m3_sm107_dense_qk_norm_rope_fp8_fake) +from .kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py new file mode 100644 index 000000000000..44365a74ba2b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared MiniMax-M3 kernel support.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py new file mode 100644 index 000000000000..8e589f83d1c0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake implementations of the MiniMax-M3 decode ops for compile and fake-tensor tracing. + +Target packages register these implementations when imported after routing. +""" + +import torch + + +@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 + + +# The MiniMax-M3 sm_107 TP2 / TP4 / TP8 decode ops: one producer and one all-reduce + norm op per TP width (Q and +# index-Q heads per rank from the kernels' producerPw{2,4,8}.h), and the dense QK-norm + RoPE ops. +def _register_minimax_m3_sm107_fakes(tp: int, heads_q: int, heads_index: int) -> None: + @torch.library.register_fake(f"trtllm::minimax_m3_sm107_tp{tp}_producer_pw") + def _( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + prewait_step_inputs=False, + ): + # [Q, index Q] in FP8; the K / V / index-K rows go into the caches in place. + num_tokens = packed.shape[0] + return [ + packed.new_empty((num_tokens, heads_q, 128), dtype=torch.float8_e4m3fn), + packed.new_empty((num_tokens, heads_index, 128), dtype=torch.float8_e4m3fn), + ] + + @torch.library.register_fake(f"trtllm::minimax_m3_sm107_tp{tp}_mc_allreduce_norm") + def _( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + prepushed, + opts, + cap_out=None, + cap_w=None, + cap_eps=0.0, + ): + # [normed, updated residual] (+ [E4M3 rows, swizzled scale bytes] with quant); cap_out is written in place. + outputs = [torch.empty_like(residual), torch.empty_like(residual)] + if quant: + m, h = residual.shape + outputs += [ + residual.new_empty((m, h), dtype=torch.float8_e4m3fn), + residual.new_empty(((m + 127) // 128 * 128 * (h // 32),), dtype=torch.uint8), + ] + return outputs + + +for _tp, _heads_q, _heads_index in ((2, 32, 2), (4, 16, 1), (8, 8, 1)): + _register_minimax_m3_sm107_fakes(_tp, _heads_q, _heads_index) + + +def _minimax_m3_sm107_dense_qk_norm_rope_fp8_fake(qkv, *args, **kwargs): + # A new FP8 tensor shaped like qkv; the KV variant also stores K / V (and the compact Q) in place. + return torch.empty_like(qkv, dtype=torch.float8_e4m3fn) + + +for _name in ( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8", + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv", + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8", + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv", +): + torch.library.register_fake(f"trtllm::{_name}")(_minimax_m3_sm107_dense_qk_norm_rope_fp8_fake) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/__init__.py index 33b2434f5049..4550fe8e4de7 100644 --- 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 @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_103 / tp4.""" -from .. import fake_ops # noqa: F401 +from ..kernels import fake_ops # noqa: F401 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 index f7e603e6af6f..91238370999c 100644 --- 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 @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_103 / tp8.""" -from .. import fake_ops # noqa: F401 +from ..kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py index d1979a7761d7..61f0a4bad521 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_107 / tp2.""" -from .. import fake_ops # noqa: F401 +from ..kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py index cca058a4265e..302bb6b15c31 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_107 / tp4.""" -from .. import fake_ops # noqa: F401 +from ..kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py index a2bc9cac05c9..856847d0b0fb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_107 / tp8.""" -from .. import fake_ops # noqa: F401 +from ..kernels import fake_ops # noqa: F401 From 53d48ee1f950a38f8a6c9a6f8181476b0065f2ae Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:35:01 -0700 Subject: [PATCH 28/29] [None][fix] Keep shared M3 fake ops in the kernel package Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../modeling_v2/models/minimax_m3/fake_ops.py | 94 +----------------- .../models/minimax_m3/kernels/__init__.py | 3 + .../models/minimax_m3/kernels/fake_ops.py | 95 +++++++++++++++++++ .../minimax_m3_nvfp4__sm_103__tp4/__init__.py | 2 +- .../minimax_m3_nvfp4__sm_103__tp8/__init__.py | 2 +- 5 files changed, 102 insertions(+), 94 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py 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 index a031ea7f9e70..76ea86dfc52e 100644 --- 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 @@ -1,95 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Fake implementations of the MiniMax-M3 decode ops for compile and fake-tensor tracing. +"""Register MiniMax-M3 fake ops through their shared kernel module.""" -Target packages register these implementations when imported after routing. -""" - -import torch - - -@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 +from .kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py new file mode 100644 index 000000000000..44365a74ba2b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared MiniMax-M3 kernel support.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py new file mode 100644 index 000000000000..a031ea7f9e70 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake implementations of the MiniMax-M3 decode ops for compile and fake-tensor tracing. + +Target packages register these implementations when imported after routing. +""" + +import torch + + +@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/_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 index 33b2434f5049..4550fe8e4de7 100644 --- 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 @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_103 / tp4.""" -from .. import fake_ops # noqa: F401 +from ..kernels import fake_ops # noqa: F401 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 index f7e603e6af6f..91238370999c 100644 --- 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 @@ -2,4 +2,4 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax-M3 NVFP4 / sm_103 / tp8.""" -from .. import fake_ops # noqa: F401 +from ..kernels import fake_ops # noqa: F401 From 74002aee2af350ee7f50d9217aa10a036772a447 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:35:58 -0700 Subject: [PATCH 29/29] [None][fix] Align graph fixtures and import step-input fake registration Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tests/unittest/_torch/executor/test_cuda_graph_runner.py | 3 --- .../_torch/kernels/parallel_hw_agnostic/test_custom_ops.py | 1 + 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unittest/_torch/executor/test_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_cuda_graph_runner.py index b26208a1ba7e..dc443ec2d7cc 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_runner.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_runner.py @@ -16,7 +16,6 @@ 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 @@ -222,13 +221,11 @@ def _decoder_runner( max_num_tokens=2, use_mrope=False, spec_config=None, - cuda_graph_mem_pool=None, enable_attention_dp=False, 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, diff --git a/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py b/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py index 0c9a6a6a5010..c39ee3ada042 100644 --- a/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py +++ b/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py @@ -21,6 +21,7 @@ # The MiniMax-M3 ops' fake impls are registered by their modeling_v2 package. import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.fake_ops # noqa: F401 +import tensorrt_llm._torch.pyexecutor.step_inputs # noqa: F401 from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops import ( _cpp_fakes, # noqa: F401 )